include llama.cpp since there were some edits

This commit is contained in:
zefie committed 2026-11-01 16:32:35 -05:00
1 parent db7f28d087
commit 8c66990853
3918 files changed
+1640170 -1

No files matched your search

+2
View File
@@ -0,0 +1,2 @@
.venv
tmp
+96
View File
@@ -0,0 +1,96 @@
# Server tests
Python based server tests scenario using [pytest](https://docs.pytest.org/en/stable/).
Tests target GitHub workflows job runners with 4 vCPU.
Note: If the host architecture inference speed is faster than GitHub runners one, parallel scenario may randomly fail.
To mitigate it, you can increase values in `n_predict`, `kv_size`.
### Install dependencies
`pip install -r requirements.txt`
### Run tests
1. Build the server
```shell
cd ../../..
cmake -B build
cmake --build build --target llama-server
```
2. Start the test: `./tests.sh`
It's possible to override some scenario steps values with environment variables:
| variable | description |
|--------------------------|------------------------------------------------------------------------------------------------|
| `PORT` | `context.server_port` to set the listening port of the server during scenario, default: `8080` |
| `LLAMA_SERVER_BIN_PATH` | to change the server binary path, default: `../../../build/bin/llama-server` |
| `DEBUG` | to enable steps and server verbose mode `--verbose` |
| `N_GPU_LAYERS` | number of model layers to offload to VRAM `-ngl --n-gpu-layers` |
| `LLAMA_CACHE` | by default server tests re-download models to the `tmp` subfolder. Set this to your cache (e.g. `$HOME/Library/Caches/llama.cpp` on Mac or `$HOME/.cache/llama.cpp` on Unix) to avoid this |
To run slow tests (will download many models, make sure to set `LLAMA_CACHE` if needed):
```shell
SLOW_TESTS=1 ./tests.sh
```
To run with stdout/stderr display in real time (verbose output, but useful for debugging):
```shell
DEBUG=1 ./tests.sh -s -v -x
```
To run all the tests in a file:
```shell
./tests.sh unit/test_chat_completion.py -v -x
```
To run a single test:
```shell
./tests.sh unit/test_chat_completion.py::test_invalid_chat_completion_req
```
Hint: You can compile and run test in single command, useful for local development:
```shell
cmake --build build -j --target llama-server && ./tools/server/tests/tests.sh
```
To see all available arguments, please refer to [pytest documentation](https://docs.pytest.org/en/stable/how-to/usage.html)
### Debugging external llama-server
It can sometimes be useful to run the server in a debugger when invesigating test
failures. To do this, the environment variable `DEBUG_EXTERNAL=1` can be set
which will cause the test to skip starting a llama-server itself. Instead, the
server can be started in a debugger.
Example using `gdb`:
```console
$ gdb --args ../../../build/bin/llama-server \
--host 127.0.0.1 --port 8080 \
--temp 0.8 --seed 42 \
--hf-repo ggml-org/models --hf-file tinyllamas/stories260K.gguf \
--batch-size 32 --no-slots --alias tinyllama-2 --ctx-size 512 \
--parallel 2 --n-predict 64
```
And a break point can be set in before running:
```console
(gdb) br server.cpp:4604
(gdb) r
main: server is listening on http://127.0.0.1:8080 - starting the main loop
srv update_slots: all slots are idle
```
And then the test in question can be run in another terminal:
```console
(venv) $ env DEBUG_EXTERNAL=1 ./tests.sh unit/test_chat_completion.py -v -x
```
And this should trigger the breakpoint and allow inspection of the server state
in the debugger terminal.
+35
View File
@@ -0,0 +1,35 @@
import os
import pytest
from filelock import FileLock
from utils import *
@pytest.fixture(scope="session", autouse=True)
def configure_worker_port(request):
worker_id = getattr(request.config, "workerinput", {}).get("workerid", "master")
worker_num = 0 if worker_id == "master" else int(worker_id[2:])
port_base = int(os.environ.get("LLAMA_SERVER_TEST_PORT_BASE", "8080"))
os.environ["PORT"] = str(port_base + worker_num * 10)
# ref: https://stackoverflow.com/questions/22627659/run-code-before-and-after-each-test-in-py-test
@pytest.fixture(autouse=True)
def stop_server_after_each_test():
# do nothing before each test
yield
# stop all servers after each test
instances = set(
server_instances
) # copy the set to prevent 'Set changed size during iteration'
for server in instances:
server.stop()
@pytest.fixture(scope="session", autouse=True)
def load_server_presets(configure_worker_port, tmp_path_factory):
# this will be run once per test session, before any tests
# serialize model downloads across parallel workers.
root_tmp_dir = tmp_path_factory.getbasetemp().parent
with FileLock(str(root_tmp_dir / "load_all.lock")):
ServerPreset.load_all()
@@ -0,0 +1,118 @@
#!/usr/bin/env python3
"""
Minimal MCP server that writes notification + response in a single write() with no flush.
This reproduces the buffering bug where read_message() can strand the response.
"""
import json
import sys
import os
TOOLS = [
{
"name": "echo",
"description": "Echo back the input message",
"inputSchema": {
"type": "object",
"properties": {
"message": {"type": "string"}
},
"required": ["message"]
}
}
]
def handle_initialize(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"protocolVersion": "2024-11-05",
"capabilities": {"tools": {}},
"serverInfo": {"name": "burst-test", "version": "1.0"}
}
}
def handle_tools_list(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {"tools": TOOLS}
}
def handle_tools_call(params, req_id):
tool_name = params.get("name")
arguments = params.get("arguments", {})
if tool_name == "echo":
message = arguments.get("message", "")
notif = {
"jsonrpc": "2.0",
"method": "notifications/progress",
"params": {"progress": 50, "total": 100}
}
response = {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"content": [{"type": "text", "text": f"echo: {message}"}]
}
}
# Single os.write() call: both lines land in one pipe packet atomically.
# This is the key difference from mcp_malformed_server.py which flushes between writes.
data = (json.dumps(notif) + "\n" + json.dumps(response) + "\n").encode("utf-8")
os.write(sys.stdout.fileno(), data)
return None # already written
else:
response = {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
}
return response
HANDLERS = {
"initialize": handle_initialize,
"tools/list": handle_tools_list,
"tools/call": handle_tools_call,
}
def main():
# Use line-buffered text mode for regular responses, but the burst write
# uses os.write() directly to guarantee a single kernel write().
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
for line in sys.stdin:
line = line.strip()
if not line:
continue
try:
request = json.loads(line)
except json.JSONDecodeError:
continue
method = request.get("method")
req_id = request.get("id")
params = request.get("params", {})
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
if req_id is None:
continue
handler = HANDLERS.get(method)
if handler:
response = handler(params, req_id)
if response is not None:
sys.stdout.write(json.dumps(response) + "\n")
sys.stdout.flush()
else:
response = {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32601, "message": f"Method not found: {method}"}
}
sys.stdout.write(json.dumps(response) + "\n")
sys.stdout.flush()
if __name__ == "__main__":
main()
@@ -0,0 +1,114 @@
#!/usr/bin/env python3
"""
MCP server that crashes after receiving a specific tool call.
"""
import json
import sys
import os
def handle_initialize(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"protocolVersion": "2024-11-05",
"capabilities": {"tools": {}},
"serverInfo": {"name": "crash-test", "version": "1.0"}
}
}
def handle_tools_list(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"tools": [
{
"name": "echo",
"description": "Echo back the input message",
"inputSchema": {
"type": "object",
"properties": {
"message": {"type": "string"}
}
}
},
{
"name": "crash",
"description": "Crash the server",
"inputSchema": {
"type": "object",
"properties": {}
}
}
]
}
}
def handle_tools_call(params, req_id):
tool_name = params.get("name")
arguments = params.get("arguments", {})
if tool_name == "echo":
message = arguments.get("message", "")
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"content": [{"type": "text", "text": f"echo: {message}"}]
}
}
elif tool_name == "crash":
# Send a partial response then exit
sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": "crashing..."}]}}) + "\n")
sys.stdout.flush()
os._exit(1)
else:
return {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
}
HANDLERS = {
"initialize": handle_initialize,
"tools/list": handle_tools_list,
"tools/call": handle_tools_call,
}
def main():
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
for line in sys.stdin:
line = line.strip()
if not line:
continue
try:
request = json.loads(line)
except json.JSONDecodeError:
continue
method = request.get("method")
req_id = request.get("id")
params = request.get("params", {})
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
if req_id is None:
continue
handler = HANDLERS.get(method)
if handler:
response = handler(params, req_id)
else:
response = {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32601, "message": f"Method not found: {method}"}
}
sys.stdout.write(json.dumps(response) + "\n")
sys.stdout.flush()
if __name__ == "__main__":
main()
+164
View File
@@ -0,0 +1,164 @@
#!/usr/bin/env python3
"""
Minimal MCP server for testing.
Implements JSON-RPC 2.0 over stdio (line-delimited JSON).
"""
import json
import sys
import os
# Ensure we use python3 from the current environment
if sys.platform == "win32":
# On Windows, we need to use the same python interpreter
pass
TOOLS = [
{
"name": "echo",
"description": "Echo back the input message",
"inputSchema": {
"type": "object",
"properties": {
"message": {"type": "string", "description": "Message to echo"}
},
"required": ["message"]
}
},
{
"name": "add",
"description": "Add two numbers",
"inputSchema": {
"type": "object",
"properties": {
"a": {"type": "number"},
"b": {"type": "number"}
},
"required": ["a", "b"]
}
},
{
"name": "fail_once",
"description": "Fails on first call, succeeds on subsequent calls",
"inputSchema": {
"type": "object",
"properties": {}
}
}
]
_state = {"fail_once_called": False}
def handle_initialize(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"protocolVersion": "2024-11-05",
"capabilities": {"tools": {}},
"serverInfo": {"name": "echo-test", "version": "1.0"}
}
}
def handle_tools_list(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {"tools": TOOLS}
}
def handle_tools_call(params, req_id):
tool_name = params.get("name")
arguments = params.get("arguments", {})
if tool_name == "echo":
message = arguments.get("message", "")
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"content": [{"type": "text", "text": f"echo: {message}"}]
}
}
elif tool_name == "add":
a = arguments.get("a", 0)
b = arguments.get("b", 0)
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"content": [{"type": "text", "text": str(a + b)}]
}
}
elif tool_name == "fail_once":
if not _state["fail_once_called"]:
_state["fail_once_called"] = True
return {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32000, "message": "transient error"}
}
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"content": [{"type": "text", "text": "ok"}]
}
}
else:
return {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
}
def handle_ping(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {}
}
HANDLERS = {
"initialize": handle_initialize,
"tools/list": handle_tools_list,
"tools/call": handle_tools_call,
"ping": handle_ping,
}
def main():
# Use unbuffered output
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
for line in sys.stdin:
line = line.strip()
if not line:
continue
try:
request = json.loads(line)
except json.JSONDecodeError:
continue
method = request.get("method")
req_id = request.get("id")
params = request.get("params", {})
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
if req_id is None:
continue
handler = HANDLERS.get(method)
if handler:
response = handler(params, req_id)
else:
response = {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32601, "message": f"Method not found: {method}"}
}
sys.stdout.write(json.dumps(response) + "\n")
sys.stdout.flush()
if __name__ == "__main__":
main()
@@ -0,0 +1,100 @@
#!/usr/bin/env python3
"""
MCP server (NDJSON JSON-RPC over stdio) that spawns a long-lived grandchild which inherits
this process's stdin/stdout/stderr and keeps them open.
This reproduces the reader-teardown deadlock: killing the direct MCP child (SIGKILL, which is
all subprocess_terminate() does) does NOT close the stdout/stderr pipe write ends, because the
grandchild still holds them. A server that reads those pipes with a blocking read would then
wait forever for an EOF that never arrives, hanging teardown (both warmup shutdown at startup
and process shutdown). The polled, running-aware reader must exit regardless.
"""
import json
import os
import subprocess
import sys
# Spawn a grandchild that inherits our std handles (fds 0/1/2 = the MCP pipes) and lives well
# past any teardown in the tests. We do NOT redirect its stdio, so it keeps the pipe write ends
# open even after this process is killed.
subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
TOOLS = [
{
"name": "echo",
"description": "Echo back the input message",
"inputSchema": {
"type": "object",
"properties": {"message": {"type": "string", "description": "Message to echo"}},
"required": ["message"],
},
}
]
def handle_initialize(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"protocolVersion": "2024-11-05",
"capabilities": {"tools": {}},
"serverInfo": {"name": "grandchild-test", "version": "1.0"},
},
}
def handle_tools_list(params, req_id):
return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}}
def handle_tools_call(params, req_id):
if params.get("name") == "echo":
message = params.get("arguments", {}).get("message", "")
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {"content": [{"type": "text", "text": f"echo: {message}"}]},
}
return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32602, "message": "Unknown tool"}}
HANDLERS = {
"initialize": handle_initialize,
"tools/list": handle_tools_list,
"tools/call": handle_tools_call,
}
def main():
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
for line in sys.stdin:
line = line.strip()
if not line:
continue
try:
request = json.loads(line)
except json.JSONDecodeError:
continue
method = request.get("method")
req_id = request.get("id")
params = request.get("params", {})
if req_id is None:
continue # notification, no response
handler = HANDLERS.get(method)
if handler:
response = handler(params, req_id)
else:
response = {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Method not found: {method}"}}
sys.stdout.write(json.dumps(response) + "\n")
sys.stdout.flush()
if __name__ == "__main__":
main()
@@ -0,0 +1,113 @@
#!/usr/bin/env python3
"""
MCP server that sends malformed responses and notifications during requests.
"""
import json
import sys
import os
def handle_initialize(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"protocolVersion": "2024-11-05",
"capabilities": {"tools": {}},
"serverInfo": {"name": "malformed-test", "version": "1.0"}
}
}
def handle_tools_list(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"tools": [
{
"name": "echo",
"description": "Echo back the input message",
"inputSchema": {
"type": "object",
"properties": {
"message": {"type": "string"}
}
}
}
]
}
}
def handle_tools_call(params, req_id):
tool_name = params.get("name")
arguments = params.get("arguments", {})
if tool_name == "echo":
message = arguments.get("message", "")
# Send a notification first (no id field)
notif = {
"jsonrpc": "2.0",
"method": "notifications/progress",
"params": {"progress": 50, "total": 100}
}
sys.stdout.write(json.dumps(notif) + "\n")
sys.stdout.flush()
# Then send the actual response
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"content": [{"type": "text", "text": f"echo: {message}"}]
}
}
else:
return {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
}
HANDLERS = {
"initialize": handle_initialize,
"tools/list": handle_tools_list,
"tools/call": handle_tools_call,
}
def main():
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
for line in sys.stdin:
line = line.strip()
if not line:
continue
try:
request = json.loads(line)
except json.JSONDecodeError:
# Send malformed JSON response
sys.stdout.write("THIS IS NOT JSON\n")
sys.stdout.flush()
continue
method = request.get("method")
req_id = request.get("id")
params = request.get("params", {})
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
if req_id is None:
continue
handler = HANDLERS.get(method)
if handler:
response = handler(params, req_id)
else:
response = {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32601, "message": f"Method not found: {method}"}
}
sys.stdout.write(json.dumps(response) + "\n")
sys.stdout.flush()
if __name__ == "__main__":
main()
+132
View File
@@ -0,0 +1,132 @@
#!/usr/bin/env python3
"""
MCP server that sleeps before responding, for timeout testing.
"""
import json
import sys
import os
import time
import argparse
TOOLS = [
{
"name": "sleep",
"description": "Sleep for a given number of seconds",
"inputSchema": {
"type": "object",
"properties": {
"seconds": {"type": "number", "description": "Seconds to sleep"}
},
"required": ["seconds"]
}
}
]
def handle_initialize(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"protocolVersion": "2024-11-05",
"capabilities": {"tools": {}},
"serverInfo": {"name": "slow-test", "version": "1.0"}
}
}
def handle_tools_list(params, req_id):
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {"tools": TOOLS}
}
def handle_tools_call(params, req_id):
tool_name = params.get("name")
arguments = params.get("arguments", {})
if tool_name == "sleep":
seconds = arguments.get("seconds", 1)
time.sleep(seconds)
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"content": [{"type": "text", "text": f"slept {seconds}s"}]
}
}
else:
return {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
}
HANDLERS = {
"initialize": handle_initialize,
"tools/list": handle_tools_list,
"tools/call": handle_tools_call,
}
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--delay", type=float, default=5.0, help="Delay in seconds for sleep tool")
args = parser.parse_args()
# Override the sleep duration
global handle_tools_call
def handle_tools_call(params, req_id):
tool_name = params.get("name")
arguments = params.get("arguments", {})
if tool_name == "sleep":
seconds = arguments.get("seconds", args.delay)
time.sleep(seconds)
return {
"jsonrpc": "2.0",
"id": req_id,
"result": {
"content": [{"type": "text", "text": f"slept {seconds}s"}]
}
}
else:
return {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
}
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
for line in sys.stdin:
line = line.strip()
if not line:
continue
try:
request = json.loads(line)
except json.JSONDecodeError:
continue
method = request.get("method")
req_id = request.get("id")
params = request.get("params", {})
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
if req_id is None:
continue
handler = HANDLERS.get(method)
if handler:
response = handler(params, req_id)
else:
response = {
"jsonrpc": "2.0",
"id": req_id,
"error": {"code": -32601, "message": f"Method not found: {method}"}
}
sys.stdout.write(json.dumps(response) + "\n")
sys.stdout.flush()
if __name__ == "__main__":
main()
+4
View File
@@ -0,0 +1,4 @@
[pytest]
markers =
slow: marks tests as slow (deselect with '-m "not slow"')
serial
@@ -0,0 +1,9 @@
aiohttp~=3.9.3
pytest~=8.3.3
pytest-xdist~=3.6
filelock~=3.16
numpy~=2.2.6
openai~=2.14.0
prometheus-client~=0.20.0
requests~=2.32.3
wget~=3.2
+26
View File
@@ -0,0 +1,26 @@
#!/usr/bin/env bash
# make sure we are in the right directory
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
cd $SCRIPT_DIR
set -eu
WORKERS="${PYTEST_WORKERS:-4}"
if [ "${WORKERS}" -eq 1 ]; then
WORKER_FLAGS=""
else
WORKER_FLAGS="-n ${WORKERS} --dist=worksteal"
fi
if [ $# -lt 1 ]
then
if [[ "${SLOW_TESTS:-0}" == 1 ]]; then
pytest --durations=30 -v -x ${WORKER_FLAGS}
else
pytest --durations=30 -v -x ${WORKER_FLAGS} -m "not slow"
fi
else
pytest --durations=30 ${WORKER_FLAGS} "$@"
fi
@@ -0,0 +1,145 @@
import pytest
import requests
import socket
from utils import *
server = ServerPreset.tinyllama2()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
def test_server_start_simple():
global server
server.start()
res = server.make_request("GET", "/health")
assert res.status_code == 200
def test_server_multiple_addresses(monkeypatch):
# The CLI value replaces the environment value, including an unavailable address.
monkeypatch.setenv("LLAMA_ARG_HOST", "192.0.2.1")
try:
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as probe:
probe.bind(("::1", 0))
except OSError:
pytest.skip("IPv6 loopback is unavailable") # ty: ignore[too-many-positional-arguments]
server.server_host = "127.0.0.1,::1"
server.api_key = "test-multiple-addresses"
server.start()
def check_address(host):
res = server.make_request("GET", "/health", host=host)
assert res.status_code == 200
res = server.make_request("POST", "/v1/completions", data={}, host=host)
assert res.status_code == 401
events = list(server.make_stream_request("POST", "/v1/completions", data={
"prompt": "Once upon a time",
"max_tokens": 8,
"stream": True,
}, headers={"Authorization": f"Bearer {server.api_key}"}, host=host))
assert len(events) > 1
return True
# parallel_function_calls swallows exceptions, a failed check leaves None in the results
results = parallel_function_calls([(check_address, (host,)) for host in ["127.0.0.1", "[::1]"]])
assert all(results)
def test_server_props():
global server
server.start()
res = server.make_request("GET", "/props")
assert res.status_code == 200
assert ".gguf" in res.body["model_path"]
assert res.body["total_slots"] == server.n_slots
default_val = res.body["default_generation_settings"]
assert server.n_ctx is not None and server.n_slots is not None
assert default_val["n_ctx"] == server.n_ctx / server.n_slots
assert default_val["params"]["seed"] == server.seed
def test_server_models():
global server
server.start()
res = server.make_request("GET", "/models")
assert res.status_code == 200
assert len(res.body["data"]) == 1
assert res.body["data"][0]["id"] == server.model_alias
def test_server_slots():
global server
# without slots endpoint enabled, this should return error
server.server_slots = False
server.start()
res = server.make_request("GET", "/slots")
assert res.status_code == 501 # ERROR_TYPE_NOT_SUPPORTED
assert "error" in res.body
server.stop()
# with slots endpoint enabled, this should return slots info
server.server_slots = True
server.n_slots = 2
server.start()
res = server.make_request("GET", "/slots")
assert res.status_code == 200
assert len(res.body) == server.n_slots
assert server.n_ctx is not None and server.n_slots is not None
assert res.body[0]["n_ctx"] == server.n_ctx / server.n_slots
assert "params" not in res.body[0]
def test_load_split_model():
global server
server.offline = False
server.model_hf_repo = "ggml-org/models"
server.model_hf_file = "tinyllamas/split/stories15M-q8_0-00001-of-00003.gguf"
server.model_alias = "tinyllama-split"
server.start()
res = server.make_request("POST", "/completion", data={
"n_predict": 16,
"prompt": "Hello",
"temperature": 0.0,
})
assert res.status_code == 200
assert match_regex("(little|girl)+", res.body["content"])
def test_no_ui():
global server
# default: UI enabled
server.start()
url = f"http://{server.server_host}:{server.server_port}"
res = requests.get(url)
assert res.status_code == 200
assert "<!doctype html>" in res.text
server.stop()
# with --no-ui, the UI should be disabled
server.no_ui = True
server.start()
res = requests.get(url)
assert res.status_code == 404
def test_server_model_aliases_and_tags():
global server
server.model_alias = "tinyllama-2,fim,code"
server.model_tags = "chat,fim,small"
server.start()
res = server.make_request("GET", "/models")
assert res.status_code == 200
assert len(res.body["data"]) == 1
model = res.body["data"][0]
# aliases field must contain all aliases
assert set(model["aliases"]) == {"tinyllama-2", "fim", "code"}
# tags field must contain all tags
assert set(model["tags"]) == {"chat", "fim", "small"}
# id is derived from first alias (alphabetical order from std::set)
assert model["id"] == "code"
@@ -0,0 +1,629 @@
import pytest
from openai import OpenAI
from utils import *
server: ServerProcess
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
@pytest.mark.parametrize(
"model,system_prompt,user_prompt,max_tokens,re_content,n_prompt,n_predicted,finish_reason,jinja,chat_template",
[
(None, "Book", "Hey", 8, "But she couldn't", 69, 8, "length", False, None),
(None, "Book", "Hey", 8, "But she couldn't", 69, 8, "length", True, None),
(None, "Book", "What is the best book", 8, "(Suddenly)+|\\{ \" Sarax.", 77, 8, "length", False, None),
(None, "Book", "What is the best book", 8, "(Suddenly)+|\\{ \" Sarax.", 77, 8, "length", True, None),
(None, "Book", "What is the best book", 8, "(Suddenly)+|\\{ \" Sarax.", 77, 8, "length", True, 'chatml'),
(None, "Book", "What is the best book", 8, "^ blue", 23, 8, "length", True, "This is not a chat template, it is"),
("codellama70b", "You are a coding assistant.", "Write the fibonacci function in c++.", 128, "(Aside|she|felter|alonger)+", 104, 128, "length", False, None),
("codellama70b", "You are a coding assistant.", "Write the fibonacci function in c++.", 128, "(Aside|she|felter|alonger)+", 104, 128, "length", True, None),
(None, "Book", [{"type": "text", "text": "What is"}, {"type": "text", "text": "the best book"}], 8, "Whillicter", 79, 8, "length", False, None),
(None, "Book", [{"type": "text", "text": "What is"}, {"type": "text", "text": "the best book"}], 8, "Whillicter", 79, 8, "length", True, None),
]
)
def test_chat_completion(model, system_prompt, user_prompt, max_tokens, re_content, n_prompt, n_predicted, finish_reason, jinja, chat_template):
global server
server.jinja = jinja
server.chat_template = chat_template
server.start()
res = server.make_request("POST", "/chat/completions", data={
"model": model,
"max_tokens": max_tokens,
"messages": [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
})
assert res.status_code == 200
assert "cmpl" in res.body["id"] # make sure the completion id has the expected format
assert res.body["system_fingerprint"].startswith("b")
# we no longer reflect back the model name, see https://github.com/ggml-org/llama.cpp/pull/17668
# assert res.body["model"] == model if model is not None else server.model_alias
assert res.body["usage"]["prompt_tokens"] == n_prompt
assert res.body["usage"]["completion_tokens"] == n_predicted
choice = res.body["choices"][0]
assert "assistant" == choice["message"]["role"]
assert match_regex(re_content, choice["message"]["content"]), f'Expected {re_content}, got {choice["message"]["content"]}'
assert choice["finish_reason"] == finish_reason
def test_chat_completion_cached_tokens():
global server
server.n_slots = 1
server.start()
seq = [
("1 2 3 4 5 6", 77, 0),
("1 2 3 4 5 6", 77, 76),
("1 2 3 4 5 9", 77, 51),
("1 2 3 9 9 9", 77, 47),
]
for user_prompt, n_prompt, n_cache in seq:
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": 8,
"messages": [
{"role": "system", "content": "Test"},
{"role": "user", "content": user_prompt},
],
})
assert res.body["usage"]["prompt_tokens"] == n_prompt
assert res.body["usage"]["prompt_tokens_details"]["cached_tokens"] == n_cache
@pytest.mark.parametrize(
"system_prompt,user_prompt,max_tokens,re_content,n_prompt,n_predicted,finish_reason",
[
("Book", "What is the best book", 8, "(Suddenly)+", 77, 8, "length"),
("You are a coding assistant.", "Write the fibonacci function in c++.", 128, "(Aside|she|felter|alonger)+", 104, 128, "length"),
]
)
def test_chat_completion_stream(system_prompt, user_prompt, max_tokens, re_content, n_prompt, n_predicted, finish_reason):
global server
server.model_alias = "llama-test-model"
server.start()
res = server.make_stream_request("POST", "/chat/completions", data={
"max_tokens": max_tokens,
"messages": [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
"stream": True,
})
content = ""
last_cmpl_id = None
for i, data in enumerate(res):
if data["choices"]:
choice = data["choices"][0]
if i == 0:
# Check first role message for stream=True
assert choice["delta"]["content"] is None
assert choice["delta"]["role"] == "assistant"
else:
assert "role" not in choice["delta"]
assert data["system_fingerprint"].startswith("b")
assert data["model"] == "llama-test-model"
if last_cmpl_id is None:
last_cmpl_id = data["id"]
assert last_cmpl_id == data["id"] # make sure the completion id is the same for all events in the stream
if choice["finish_reason"] in ["stop", "length"]:
assert "content" not in choice["delta"]
assert match_regex(re_content, content)
assert choice["finish_reason"] == finish_reason
else:
assert choice["finish_reason"] is None
content += choice["delta"]["content"] or ''
else:
assert data["usage"]["prompt_tokens"] == n_prompt
assert data["usage"]["completion_tokens"] == n_predicted
def test_chat_completion_with_openai_library():
global server
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.chat.completions.create(
model="gpt-3.5-turbo-instruct",
messages=[
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
],
max_tokens=8,
seed=42,
temperature=0.8,
)
assert res.system_fingerprint is not None and res.system_fingerprint.startswith("b")
assert res.choices[0].finish_reason == "length"
assert res.choices[0].message.content is not None
assert match_regex("(Suddenly)+", res.choices[0].message.content)
def test_chat_template():
global server
server.chat_template = "llama3"
server.debug = True # to get the "__verbose" object in the response
server.start()
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": 8,
"messages": [
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
]
})
assert res.status_code == 200
assert "__verbose" in res.body
assert res.body["__verbose"]["prompt"] == "<s> <|start_header_id|>system<|end_header_id|>\n\nBook<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nWhat is the best book<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
@pytest.mark.parametrize("prefill,re_prefill", [
("Whill", "Whill"),
([{"type": "text", "text": "Wh"}, {"type": "text", "text": "ill"}], "Wh\n\nill"),
])
def test_chat_template_assistant_prefill(prefill, re_prefill):
global server
server.jinja = True
server.chat_template_file = "../../../models/templates/meta-llama-Llama-3.1-8B-Instruct.jinja"
server.debug = True # to get the "__verbose" object in the response
server.start()
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": 8,
"messages": [
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
{"role": "assistant", "content": prefill},
]
})
assert res.status_code == 200
assert "__verbose" in res.body
assert res.body["__verbose"]["prompt"].endswith(f"<|start_header_id|>user<|end_header_id|>\n\nWhat is the best book<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n{re_prefill}")
def test_chat_template_continue_final_message_vllm_compat():
"""continue_final_message is the vLLM/transformers explicit alias for the prefill_assistant heuristic.
Both must produce the same prompt."""
global server
server.jinja = True
server.chat_template_file = "../../../models/templates/meta-llama-Llama-3.1-8B-Instruct.jinja"
server.debug = True
server.start()
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": 8,
"add_generation_prompt": False,
"continue_final_message": True,
"messages": [
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
{"role": "assistant", "content": "Whill"},
]
})
assert res.status_code == 200
assert "__verbose" in res.body
assert res.body["__verbose"]["prompt"].endswith("<|start_header_id|>user<|end_header_id|>\n\nWhat is the best book<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\nWhill")
def test_chat_template_continue_final_message_mutual_exclusion():
"""add_generation_prompt and continue_final_message both set to true must be rejected"""
global server
server.chat_template = "llama3"
server.start()
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": 8,
"add_generation_prompt": True,
"continue_final_message": True,
"messages": [
{"role": "user", "content": "Hi"},
{"role": "assistant", "content": "Hello"},
]
})
assert res.status_code == 400
def test_apply_chat_template():
global server
server.chat_template = "command-r"
server.start()
res = server.make_request("POST", "/apply-template", data={
"messages": [
{"role": "system", "content": "You are a test."},
{"role": "user", "content":"Hi there"},
]
})
assert res.status_code == 200
assert "prompt" in res.body
assert res.body["prompt"] == "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>You are a test.<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|USER_TOKEN|>Hi there<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>"
@pytest.mark.parametrize("response_format,n_predicted,re_content", [
({"type": "json_object", "schema": {"const": "42"}}, 6, "\"42\""),
({"type": "json_object", "schema": {"items": [{"type": "integer"}]}}, 10, "[ -3000 ]"),
({"type": "json_schema", "json_schema": {"schema": {"const": "foooooo"}}}, 10, "\"foooooo\""),
({"type": "json_object"}, 10, "(\\{|John)+"),
({"type": "sound"}, 0, None),
# invalid response format (expected to fail)
({"type": "json_object", "schema": 123}, 0, None),
({"type": "json_object", "schema": {"type": 123}}, 0, None),
({"type": "json_object", "schema": {"type": "hiccup"}}, 0, None),
])
def test_completion_with_response_format(response_format: dict, n_predicted: int, re_content: str | None):
global server
server.start()
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": n_predicted,
"messages": [
{"role": "system", "content": "You are a coding assistant."},
{"role": "user", "content": "Write an example"},
],
"response_format": response_format,
})
if re_content is not None:
assert res.status_code == 200
choice = res.body["choices"][0]
assert match_regex(re_content, choice["message"]["content"])
else:
assert res.status_code == 400
assert "error" in res.body
@pytest.mark.parametrize("jinja,json_schema,n_predicted,re_content", [
(False, {"const": "42"}, 6, "\"42\""),
(True, {"const": "42"}, 6, "\"42\""),
])
def test_completion_with_json_schema(jinja: bool, json_schema: dict, n_predicted: int, re_content: str):
global server
server.jinja = jinja
server.debug = True
server.start()
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": n_predicted,
"messages": [
{"role": "system", "content": "You are a coding assistant."},
{"role": "user", "content": "Write an example"},
],
"json_schema": json_schema,
})
assert res.status_code == 200, f'Expected 200, got {res.status_code}'
choice = res.body["choices"][0]
assert match_regex(re_content, choice["message"]["content"]), f'Expected {re_content}, got {choice["message"]["content"]}'
@pytest.mark.parametrize("jinja,grammar,n_predicted,re_content", [
(False, 'root ::= "a"{5,5}', 6, "a{5,5}"),
(True, 'root ::= "a"{5,5}', 6, "a{5,5}"),
])
def test_completion_with_grammar(jinja: bool, grammar: str, n_predicted: int, re_content: str):
global server
server.jinja = jinja
server.start()
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": n_predicted,
"messages": [
{"role": "user", "content": "Does not matter what I say, does it?"},
],
"grammar": grammar,
})
assert res.status_code == 200, res.body
choice = res.body["choices"][0]
assert match_regex(re_content, choice["message"]["content"]), choice["message"]["content"]
def test_completion_with_invalid_grammar():
global server
server.start()
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": 8,
"messages": [
{"role": "user", "content": "Does not matter what I say, does it?"},
],
"grammar": "root ::= this is (not valid GBNF",
})
assert res.status_code == 400, res.body
assert "error" in res.body
@pytest.mark.parametrize("messages", [
None,
"string",
[123],
[{}],
[{"role": 123}],
[{"role": "system", "content": 123}],
# [{"content": "hello"}], # TODO: should not be a valid case
[{"role": "system", "content": "test"}, {}],
[{"role": "user", "content": "test"}, {"role": "assistant", "content": "test"}, {"role": "assistant", "content": "test"}],
])
def test_invalid_chat_completion_req(messages):
global server
server.start()
res = server.make_request("POST", "/chat/completions", data={
"messages": messages,
})
assert res.status_code == 400 or res.status_code == 500
assert "error" in res.body
def test_chat_completion_with_timings_per_token():
global server
server.start()
res = server.make_stream_request("POST", "/chat/completions", data={
"max_tokens": 10,
"messages": [{"role": "user", "content": "test"}],
"stream": True,
"stream_options": {"include_usage": True},
"timings_per_token": True,
})
stats_received = False
for i, data in enumerate(res):
if i == 0:
# Check first role message for stream=True
assert data["choices"][0]["delta"]["content"] is None
assert data["choices"][0]["delta"]["role"] == "assistant"
assert "timings" not in data, f'First event should not have timings: {data}'
else:
if data["choices"]:
assert "role" not in data["choices"][0]["delta"]
else:
assert "timings" in data
assert "prompt_per_second" in data["timings"]
assert "predicted_per_second" in data["timings"]
assert "predicted_n" in data["timings"]
assert data["timings"]["predicted_n"] <= 10
stats_received = True
assert stats_received
def test_logprobs():
global server
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.chat.completions.create(
model="gpt-3.5-turbo-instruct",
temperature=0.0,
messages=[
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
],
max_tokens=5,
logprobs=True,
top_logprobs=10,
)
output_text = res.choices[0].message.content
aggregated_text = ''
assert res.choices[0].logprobs is not None
assert res.choices[0].logprobs.content is not None
for token in res.choices[0].logprobs.content:
aggregated_text += token.token
assert token.logprob <= 0.0
assert token.bytes is not None
assert len(token.top_logprobs) > 0
assert aggregated_text == output_text
def test_logprobs_stream():
global server
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.chat.completions.create(
model="gpt-3.5-turbo-instruct",
temperature=0.0,
messages=[
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
],
max_tokens=5,
logprobs=True,
top_logprobs=10,
stream=True,
)
output_text = ''
aggregated_text = ''
for i, data in enumerate(res):
if data.choices:
choice = data.choices[0]
if i == 0:
# Check first role message for stream=True
assert choice.delta.content is None
assert choice.delta.role == "assistant"
else:
assert choice.delta.role is None
if choice.finish_reason is None:
if choice.delta.content:
output_text += choice.delta.content
assert choice.logprobs is not None
assert choice.logprobs.content is not None
for token in choice.logprobs.content:
aggregated_text += token.token
assert token.logprob <= 0.0
assert token.bytes is not None
assert token.top_logprobs is not None
assert len(token.top_logprobs) > 0
assert aggregated_text == output_text
def test_logit_bias():
global server
server.start()
exclude = ["i", "I", "the", "The", "to", "a", "an", "be", "is", "was", "but", "But", "and", "And", "so", "So", "you", "You", "he", "He", "she", "She", "we", "We", "they", "They", "it", "It", "his", "His", "her", "Her", "book", "Book"]
res = server.make_request("POST", "/tokenize", data={
"content": " " + " ".join(exclude) + " ",
})
assert res.status_code == 200
tokens = res.body["tokens"]
logit_bias = {tok: -100 for tok in tokens}
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.chat.completions.create(
model="gpt-3.5-turbo-instruct",
temperature=0.0,
messages=[
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
],
max_tokens=64,
logit_bias=logit_bias
)
output_text = res.choices[0].message.content
assert output_text
assert all(output_text.find(" " + tok + " ") == -1 for tok in exclude)
def test_context_size_exceeded():
global server
server.start()
res = server.make_request("POST", "/chat/completions", data={
"messages": [
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
] * 100, # make the prompt too long
})
assert res.status_code == 400
assert "error" in res.body
assert res.body["error"]["type"] == "exceed_context_size_error"
assert res.body["error"]["n_prompt_tokens"] > 0
assert server.n_ctx is not None
assert server.n_slots is not None
assert res.body["error"]["n_ctx"] == server.n_ctx // server.n_slots
def test_context_size_exceeded_stream():
global server
server.start()
try:
for _ in server.make_stream_request("POST", "/chat/completions", data={
"messages": [
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
] * 100, # make the prompt too long
"stream": True}):
pass
assert False, "Should have failed"
except ServerError as e:
assert e.code == 400
assert "error" in e.body
assert e.body["error"]["type"] == "exceed_context_size_error"
assert e.body["error"]["n_prompt_tokens"] > 0
assert server.n_ctx is not None
assert server.n_slots is not None
assert e.body["error"]["n_ctx"] == server.n_ctx // server.n_slots
@pytest.mark.parametrize(
"n_batch,batch_count,reuse_cache",
[
(64, 4, False),
(64, 2, True),
]
)
def test_return_progress(n_batch, batch_count, reuse_cache):
global server
server.n_batch = n_batch
server.n_ctx = 256
server.n_slots = 1
server.start()
def make_cmpl_request():
return server.make_stream_request("POST", "/chat/completions", data={
"max_tokens": 10,
"messages": [
{"role": "user", "content": "This is a test" * 10},
],
"stream": True,
"return_progress": True,
})
if reuse_cache:
# make a first request to populate the cache
res0 = make_cmpl_request()
for _ in res0:
pass # discard the output
res = make_cmpl_request()
last_progress = None
total_batch_count = 0
for data in res:
cur_progress = data.get("prompt_progress", None)
if cur_progress is None:
continue
if total_batch_count == 0:
# first progress report must have n_cache == n_processed
assert cur_progress["total"] > 0
assert cur_progress["cache"] == cur_progress["processed"]
if reuse_cache:
# when reusing cache, we expect some cached tokens
assert cur_progress["cache"] > 0
if last_progress is not None:
assert cur_progress["total"] == last_progress["total"]
assert cur_progress["cache"] == last_progress["cache"]
assert cur_progress["processed"] > last_progress["processed"]
total_batch_count += 1
last_progress = cur_progress
# last progress should indicate completion (all tokens processed)
assert last_progress is not None
assert last_progress["total"] > 0
assert last_progress["processed"] == last_progress["total"]
assert total_batch_count == batch_count
def test_chat_completions_multiple_choices():
global server
server.start()
# make sure cache can be reused across multiple choices and multiple requests
# ref: https://github.com/ggml-org/llama.cpp/pull/18663
for _ in range(2):
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": 8,
"n": 2,
"messages": [
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
],
# test forcing the same slot to be used
# the scheduler should not be locked up in this case
"id_slot": 0,
})
assert res.status_code == 200
assert len(res.body["choices"]) == 2
for choice in res.body["choices"]:
assert "assistant" == choice["message"]["role"]
assert choice["finish_reason"] == "length"
# each choice generates max_tokens tokens, usage counts them all
# ref: https://github.com/ggml-org/llama.cpp/issues/29451
assert res.body["usage"]["completion_tokens"] == 16
assert res.body["usage"]["total_tokens"] == res.body["usage"]["prompt_tokens"] + 16
def test_chat_completions_token_count():
global server
server.start()
# make sure cache can be reused across multiple choices and multiple requests
# ref: https://github.com/ggml-org/llama.cpp/pull/18663
for _ in range(2):
res = server.make_request("POST", "/chat/completions/input_tokens", data={
"messages": [
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
],
})
assert res.status_code == 200
assert res.body["input_tokens"] > 5
def test_verbose_debug():
global server
server.start()
for verbose in [True, False]:
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": 2,
"messages": [
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
],
"verbose": verbose,
})
assert res.status_code == 200
if verbose:
assert "__verbose" in res.body
assert "Book" in res.body["__verbose"]["prompt"]
else:
assert "__verbose" not in res.body
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,60 @@
import pytest
from utils import *
server: ServerProcess
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
server.gcp_compat = True
def test_gcp_predict_camel_case():
global server
server.start()
res = server.make_request("POST", "/predict", data={
"instances": [
{
"@requestFormat": "chatCompletions",
"max_tokens": 8,
"messages": [
{"role": "user", "content": "What is the meaning of life?"},
],
}
],
})
assert res.status_code == 200
assert "predictions" in res.body
assert len(res.body["predictions"]) == 1
prediction = res.body["predictions"][0]
assert "choices" in prediction
assert len(prediction["choices"]) == 1
assert prediction["choices"][0]["message"]["role"] == "assistant"
assert len(prediction["choices"][0]["message"]["content"]) > 0
def test_gcp_predict_multiple_instances():
global server
server.n_slots = 2
server.start()
res = server.make_request("POST", "/predict", data={
"instances": [
{
"@requestFormat": "chatCompletions",
"max_tokens": 8,
"messages": [{"role": "user", "content": "Say hello"}],
},
{
"@requestFormat": "chatCompletions",
"max_tokens": 8,
"messages": [{"role": "user", "content": "Say world"}],
},
],
})
assert res.status_code == 200
assert len(res.body["predictions"]) == 2
for prediction in res.body["predictions"]:
assert "choices" in prediction
assert len(prediction["choices"][0]["message"]["content"]) > 0
@@ -0,0 +1,114 @@
import pytest
from openai import OpenAI
from utils import *
server: ServerProcess
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
def test_responses_with_openai_library():
global server
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.responses.create(
model="gpt-4.1",
input=[
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
],
max_output_tokens=8,
temperature=0.8,
)
assert res.id.startswith("resp_")
assert res.output[0].id is not None
assert res.output[0].id.startswith("msg_")
assert match_regex("(Suddenly)+", res.output_text)
def test_responses_stream_with_openai_library():
global server
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
stream = client.responses.create(
model="gpt-4.1",
input=[
{"role": "system", "content": "Book"},
{"role": "user", "content": "What is the best book"},
],
max_output_tokens=8,
temperature=0.8,
stream=True,
)
gathered_text = ''
resp_id = ''
msg_id = ''
for r in stream:
if r.type == "response.created":
assert r.response.id.startswith("resp_")
resp_id = r.response.id
if r.type == "response.in_progress":
assert r.response.id == resp_id
if r.type == "response.output_item.added":
assert r.item.id is not None
assert r.item.id.startswith("msg_")
msg_id = r.item.id
if (r.type == "response.content_part.added" or
r.type == "response.output_text.delta" or
r.type == "response.output_text.done" or
r.type == "response.content_part.done"):
assert r.item_id == msg_id
if r.type == "response.output_item.done":
assert r.item.id == msg_id
if r.type == "response.output_text.delta":
gathered_text += r.delta
if r.type == "response.completed":
assert r.response.id.startswith("resp_")
assert r.response.output[0].id is not None
assert r.response.output[0].id.startswith("msg_")
assert gathered_text == r.response.output_text
assert match_regex("(Suddenly)+", r.response.output_text)
def test_responses_stream_with_llama_telemetry():
global server
server.n_ctx = 256
server.n_batch = 32
server.n_slots = 1
server.start()
saw_progress = False
saw_delta_timings = False
completed = None
res = server.make_stream_request("POST", "/responses", data={
"input": "This is a test" * 10,
"max_output_tokens": 8,
"temperature": 0.8,
"stream": True,
"timings_per_token": True,
"return_progress": True,
})
for data in res:
if "prompt_progress" in data:
assert data["type"] == "response.in_progress"
assert data["prompt_progress"]["total"] > 0
assert data["prompt_progress"]["processed"] >= data["prompt_progress"]["cache"]
saw_progress = True
if "timings" in data:
assert "prompt_per_second" in data["timings"]
assert "predicted_per_second" in data["timings"]
if data["type"] == "response.output_text.delta":
saw_delta_timings = True
if data["type"] == "response.completed":
completed = data
assert saw_progress
assert saw_delta_timings
assert completed is not None
assert "usage" in completed["response"]
assert "timings" in completed
@@ -0,0 +1,769 @@
import pytest
import requests
import time
import random
from openai import OpenAI
from utils import *
server = ServerPreset.tinyllama2()
JSON_MULTIMODAL_KEY = "multimodal_data"
JSON_PROMPT_STRING_KEY = "prompt_string"
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
@pytest.mark.parametrize("prompt,n_predict,re_content,n_prompt,n_predicted,truncated,return_tokens", [
("I believe the meaning of life is", 8, "(going|bed)+", 18, 8, False, False),
("Write a joke about AI from a very long prompt which will not be truncated", 64, "(princesses|everyone|kids|Anna|forest)+", 46, 64, False, True),
])
def test_completion(prompt: str, n_predict: int, re_content: str, n_prompt: int, n_predicted: int, truncated: bool, return_tokens: bool):
global server
server.start()
res = server.make_request("POST", "/completion", data={
"n_predict": n_predict,
"prompt": prompt,
"return_tokens": return_tokens,
})
assert res.status_code == 200
assert res.body["timings"]["prompt_n"] == n_prompt
assert res.body["timings"]["predicted_n"] == n_predicted
assert res.body["truncated"] == truncated
assert type(res.body["has_new_line"]) == bool
assert match_regex(re_content, res.body["content"])
if return_tokens:
assert len(res.body["tokens"]) > 0
assert all(type(tok) == int for tok in res.body["tokens"])
else:
assert res.body["tokens"] == []
@pytest.mark.parametrize("prompt,n_predict,re_content,n_prompt,n_predicted,truncated", [
("I believe the meaning of life is", 8, "(going|bed)+", 18, 8, False),
("Write a joke about AI from a very long prompt which will not be truncated", 64, "(princesses|everyone|kids|Anna|forest)+", 46, 64, False),
])
def test_completion_stream(prompt: str, n_predict: int, re_content: str, n_prompt: int, n_predicted: int, truncated: bool):
global server
server.start()
res = server.make_stream_request("POST", "/completion", data={
"n_predict": n_predict,
"prompt": prompt,
"stream": True,
})
content = ""
for data in res:
assert "stop" in data and type(data["stop"]) == bool
if data["stop"]:
assert data["timings"]["prompt_n"] == n_prompt
assert data["timings"]["predicted_n"] == n_predicted
assert data["truncated"] == truncated
assert data["stop_type"] == "limit"
assert type(data["has_new_line"]) == bool
assert "generation_settings" in data
assert server.n_predict is not None
assert data["generation_settings"]["n_predict"] == min(n_predict, server.n_predict)
assert data["generation_settings"]["seed"] == server.seed
assert "adaptive_target" in data["generation_settings"]
assert "adaptive_decay" in data["generation_settings"]
assert match_regex(re_content, content)
else:
assert len(data["tokens"]) > 0
assert all(type(tok) == int for tok in data["tokens"])
content += data["content"]
def test_completion_stream_vs_non_stream():
global server
server.start()
res_stream = server.make_stream_request("POST", "/completion", data={
"n_predict": 8,
"prompt": "I believe the meaning of life is",
"stream": True,
})
res_non_stream = server.make_request("POST", "/completion", data={
"n_predict": 8,
"prompt": "I believe the meaning of life is",
})
content_stream = ""
for data in res_stream:
content_stream += data["content"]
assert content_stream == res_non_stream.body["content"]
def test_completion_with_openai_library():
global server
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.completions.create(
model="davinci-002",
prompt="I believe the meaning of life is",
max_tokens=8,
)
assert res.system_fingerprint is not None and res.system_fingerprint.startswith("b")
assert res.choices[0].finish_reason == "length"
assert res.choices[0].text is not None
assert match_regex("(going|bed)+", res.choices[0].text)
def test_completion_stream_with_openai_library():
global server
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.completions.create(
model="davinci-002",
prompt="I believe the meaning of life is",
max_tokens=8,
stream=True,
)
output_text = ''
for data in res:
choice = data.choices[0]
if choice.finish_reason is None:
assert choice.text is not None
output_text += choice.text
assert match_regex("(going|bed)+", output_text)
# Test case from https://github.com/ggml-org/llama.cpp/issues/13780
@pytest.mark.slow
def test_completion_stream_with_openai_library_stops():
global server
server.model_hf_repo = "bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M"
server.model_hf_file = None
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.completions.create(
model="davinci-002",
prompt="System: You are helpful assistant.\nAssistant:\nHey! How could I help?\nUser:\nTell me a joke.\nAssistant:\n",
stop=["User:\n", "Assistant:\n"],
max_tokens=200,
stream=True,
)
output_text = ''
for data in res:
choice = data.choices[0]
if choice.finish_reason is None:
assert choice.text is not None
output_text += choice.text
assert match_regex("Sure, here's one for[\\s\\S]*|Sure! Here's one for you[\\s\\S]*", output_text), f'Unexpected output: {output_text}'
@pytest.mark.parametrize("n_slots", [1, 2])
def test_consistent_result_same_seed(n_slots: int):
global server
server.n_slots = n_slots
server.start()
last_res = None
for _ in range(4):
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"seed": 42,
"temperature": 0.0,
"cache_prompt": False, # TODO: remove this once test_cache_vs_nocache_prompt is fixed
})
if last_res is not None:
assert res.body["content"] == last_res.body["content"]
last_res = res
@pytest.mark.parametrize("n_slots", [1, 2])
def test_different_result_different_seed(n_slots: int):
global server
server.n_slots = n_slots
server.start()
last_res = None
for seed in range(4):
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"seed": seed,
"temperature": 1.0,
"cache_prompt": False, # TODO: remove this once test_cache_vs_nocache_prompt is fixed
})
if last_res is not None:
assert res.body["content"] != last_res.body["content"]
last_res = res
# TODO figure why it don't work with temperature = 1
# @pytest.mark.parametrize("temperature", [0.0, 1.0])
@pytest.mark.parametrize("n_batch", [16, 32])
@pytest.mark.parametrize("temperature", [0.0])
def test_consistent_result_different_batch_size(n_batch: int, temperature: float):
global server
server.n_batch = n_batch
server.start()
last_res = None
for _ in range(4):
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"seed": 42,
"temperature": temperature,
"cache_prompt": False, # TODO: remove this once test_cache_vs_nocache_prompt is fixed
})
if last_res is not None:
assert res.body["content"] == last_res.body["content"]
last_res = res
@pytest.mark.skip(reason="This test fails on linux, need to be fixed")
def test_cache_vs_nocache_prompt():
global server
server.start()
res_cache = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"seed": 42,
"temperature": 1.0,
"cache_prompt": True,
})
res_no_cache = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"seed": 42,
"temperature": 1.0,
"cache_prompt": False,
})
assert res_cache.body["content"] == res_no_cache.body["content"]
def test_nocache_long_input_prompt():
global server
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is"*32,
"seed": 42,
"temperature": 1.0,
"cache_prompt": False,
})
assert res.status_code == 400
# a request pinned to a busy slot leaves the generation running on it untouched
def test_pinned_request_on_busy_slot():
global server
server.n_ctx = 4096
server.start()
story = "Once upon a time a dragon named Ember guarded a golden key in a deep cave. " * 8
def run(pin_busy: bool) -> str:
server.make_request("POST", "/completion", data={"prompt": story, "id_slot": 0, "n_predict": 4, "temperature": 0.0})
res = server.make_stream_request("POST", "/completion", data={
"prompt": "To bake bread, mix flour, water and salt, then",
"id_slot": 0, "n_predict": 1024, "ignore_eos": True, "temperature": 0.0, "stream": True,
})
content = next(res)["content"]
if pin_busy:
server.make_request("POST", "/completion", data={"prompt": story + "The dragon", "id_slot": 0, "n_predict": 4, "temperature": 0.0})
return content + "".join(chunk["content"] for chunk in res)
baseline = run(pin_busy=False)
assert run(pin_busy=True) == baseline
def test_json_prompt_no_mtmd():
global server
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": { JSON_PROMPT_STRING_KEY: "I believe the meaning of life is" },
"seed": 42,
"temperature": 1.0,
"cache_prompt": False,
})
assert res.status_code == 200
def test_json_prompt_mtm_error_when_not_supported():
global server
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": { JSON_PROMPT_STRING_KEY: "I believe the meaning of life is <__media__>", JSON_MULTIMODAL_KEY: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" },
"seed": 42,
"temperature": 1.0,
"cache_prompt": False,
})
# MTMD is disabled on this model, so this should fail.
assert res.status_code != 200
def test_completion_with_tokens_input():
global server
server.temperature = 0.0
server.start()
prompt_str = "I believe the meaning of life is"
res = server.make_request("POST", "/tokenize", data={
"content": prompt_str,
"add_special": True,
})
assert res.status_code == 200
tokens = res.body["tokens"]
# single completion
res = server.make_request("POST", "/completion", data={
"prompt": tokens,
})
assert res.status_code == 200
assert type(res.body["content"]) == str
# batch completion
res = server.make_request("POST", "/completion", data={
"prompt": [tokens, tokens],
})
assert res.status_code == 200
assert type(res.body) == list
assert len(res.body) == 2
assert res.body[0]["content"] == res.body[1]["content"]
# mixed string and tokens
res = server.make_request("POST", "/completion", data={
"prompt": [tokens, prompt_str],
})
assert res.status_code == 200
assert type(res.body) == list
assert len(res.body) == 2
assert res.body[0]["content"] == res.body[1]["content"]
# mixed JSON and tokens
res = server.make_request("POST", "/completion", data={
"prompt": [
tokens,
{
JSON_PROMPT_STRING_KEY: "I believe the meaning of life is",
},
],
})
assert res.status_code == 200
assert type(res.body) == list
assert len(res.body) == 2
assert res.body[0]["content"] == res.body[1]["content"]
# mixed string and tokens in one sequence
res = server.make_request("POST", "/completion", data={
"prompt": [1, 2, 3, 4, 5, 6, prompt_str, 7, 8, 9, 10, prompt_str],
})
assert res.status_code == 200
assert type(res.body["content"]) == str
@pytest.mark.parametrize("n_slots,n_requests", [
(1, 3),
(2, 2),
(2, 4),
(4, 2), # some slots must be idle
(4, 6),
])
def test_completion_parallel_slots(n_slots: int, n_requests: int):
global server
server.n_slots = n_slots
server.temperature = 0.0
server.start()
PROMPTS = [
("Write a very long book.", "(very|special|big)+"),
("Write another a poem.", "(small|house)+"),
("What is LLM?", "(Dad|said)+"),
("The sky is blue and I love it.", "(climb|leaf)+"),
("Write another very long music lyrics.", "(friends|step|sky)+"),
("Write a very long joke.", "(cat|Whiskers)+"),
]
def check_slots_status():
should_all_slots_busy = n_requests >= n_slots
time.sleep(0.1)
res = server.make_request("GET", "/slots")
n_busy = sum([1 for slot in res.body if slot["is_processing"]])
if should_all_slots_busy:
assert n_busy == n_slots
else:
assert n_busy <= n_slots
tasks = []
for i in range(n_requests):
prompt, re_content = PROMPTS[i % len(PROMPTS)]
tasks.append((server.make_request, ("POST", "/completion", {
"prompt": prompt,
"seed": 42,
"temperature": 1.0,
})))
tasks.append((check_slots_status, ()))
results = parallel_function_calls(tasks)
# check results
for i in range(n_requests):
prompt, re_content = PROMPTS[i % len(PROMPTS)]
res = results[i]
assert res.status_code == 200
assert type(res.body["content"]) == str
assert len(res.body["content"]) > 10
# FIXME: the result is not deterministic when using other slot than slot 0
# assert match_regex(re_content, res.body["content"])
@pytest.mark.parametrize(
"n_ctx,n_slots,n_predict_vals,expected_success",
[
(256, 4, [80, 40, 80, 80], [True, True, True, True]),
(256, 4, [70, 70, 70, 70], [False, False, False, False]),
(256, 4, [90, 90, 40, 90], [False, False, True, False]),
(256, 4, [90, 90, 40, 75], [True, True, True, True]),
],
)
def test_completion_unified(n_ctx, n_slots, n_predict_vals, expected_success):
global server
server.n_slots = n_slots
server.kv_unified = True
server.n_ctx = n_ctx
server.start()
prompt = "A"
tasks = []
for n_predict in n_predict_vals:
tasks.append((server.make_request, ("POST", "/completion", {"prompt": prompt, "n_predict": n_predict})))
results = parallel_function_calls(tasks)
for res, n_predict, expect_ok in zip(results, n_predict_vals, expected_success):
if expect_ok:
# the pool is aborted as a whole, so a request that fits on its own
# is still dropped when the slots overlap, and it says so explicitly
assert res.status_code == 200 or (
res.status_code == 500
and "context size has been exceeded" in res.body["error"]["message"].lower()
)
# note: https://github.com/ggml-org/llama.cpp/pull/18700#issuecomment-3728695581
if res.status_code == 200:
assert "content" in res.body
if "timings" in res.body:
assert res.body["timings"]["predicted_n"] == n_predict
@pytest.mark.parametrize(
"prompt,n_predict,response_fields",
[
("I believe the meaning of life is", 8, []),
("I believe the meaning of life is", 32, ["content", "generation_settings/n_predict", "prompt"]),
],
)
def test_completion_response_fields(
prompt: str, n_predict: int, response_fields: list[str]
):
global server
server.start()
res = server.make_request(
"POST",
"/completion",
data={
"n_predict": n_predict,
"prompt": prompt,
"response_fields": response_fields,
},
)
assert res.status_code == 200
assert "content" in res.body
assert len(res.body["content"])
if len(response_fields):
assert res.body["generation_settings/n_predict"] == n_predict
assert res.body["prompt"] == "<s> " + prompt
assert isinstance(res.body["content"], str)
assert len(res.body) == len(response_fields)
else:
assert len(res.body)
assert "generation_settings" in res.body
def test_n_probs():
global server
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"n_probs": 10,
"temperature": 0.0,
"n_predict": 5,
})
assert res.status_code == 200
assert "completion_probabilities" in res.body
assert len(res.body["completion_probabilities"]) == 5
for tok in res.body["completion_probabilities"]:
assert "id" in tok and tok["id"] > 0
assert "token" in tok and type(tok["token"]) == str
assert "logprob" in tok and tok["logprob"] <= 0.0
assert "bytes" in tok and type(tok["bytes"]) == list
assert len(tok["top_logprobs"]) == 10
for prob in tok["top_logprobs"]:
assert "id" in prob and prob["id"] > 0
assert "token" in prob and type(prob["token"]) == str
assert "logprob" in prob and prob["logprob"] <= 0.0
assert "bytes" in prob and type(prob["bytes"]) == list
def test_n_probs_stream():
global server
server.start()
res = server.make_stream_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"n_probs": 10,
"temperature": 0.0,
"n_predict": 5,
"stream": True,
})
for data in res:
if data["stop"] == False:
assert "completion_probabilities" in data
assert len(data["completion_probabilities"]) == 1
for tok in data["completion_probabilities"]:
assert "id" in tok and tok["id"] > 0
assert "token" in tok and type(tok["token"]) == str
assert "logprob" in tok and tok["logprob"] <= 0.0
assert "bytes" in tok and type(tok["bytes"]) == list
assert len(tok["top_logprobs"]) == 10
for prob in tok["top_logprobs"]:
assert "id" in prob and prob["id"] > 0
assert "token" in prob and type(prob["token"]) == str
assert "logprob" in prob and prob["logprob"] <= 0.0
assert "bytes" in prob and type(prob["bytes"]) == list
def test_n_probs_post_sampling():
global server
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "Today was the day. Today I would finally become a",
"n_probs": 10,
"temperature": 1.0,
"n_predict": 5,
"post_sampling_probs": True,
})
assert res.status_code == 200
assert "completion_probabilities" in res.body
assert len(res.body["completion_probabilities"]) == 5
for (i, tok) in enumerate(res.body["completion_probabilities"]):
assert "id" in tok and tok["id"] > 0
assert "token" in tok and type(tok["token"]) == str
assert "prob" in tok and 0.0 < tok["prob"] <= 1.0
assert "bytes" in tok and type(tok["bytes"]) == list
assert "top_probs" in tok and type(tok["top_probs"]) == list
for prob in tok["top_probs"]:
assert "id" in prob and prob["id"] > 0
assert "token" in prob and type(prob["token"]) == str
# 0.0 probability tokens should never be returned by the server
assert "prob" in prob and 0.0 < prob["prob"] <= 1.0
assert "bytes" in prob and type(prob["bytes"]) == list
if i == 0:
# The prompt is vague enough that we should get at least 10 possibilities
# for the first token.
assert len(tok["top_probs"]) == 10
if len(tok["top_probs"]) < 10:
# Getting less than the requested number of probabilities should only happen
# if the ones we did get already sum to 1.0.
assert sum(p["prob"] for p in tok["top_probs"]) == pytest.approx(1.0)
def test_n_probs_post_backend_sampling():
"""Verify that the same probabilities are returned with and without backend sampling."""
global server
server.backend_sampling = True
server.start()
def make_request(backend_sampling):
n_predict = 20
res = server.make_request("POST", "/completion", data={
"prompt": "The countries of Europe, in random order, are:",
"n_probs": 10,
"n_predict": n_predict,
"post_sampling_probs": True,
"seed": 4242,
"backend_sampling": backend_sampling,
})
assert res.status_code == 200
total_probs = 0
completions = res.body["completion_probabilities"]
assert len(completions) == n_predict
for tok in completions:
# Handling of 0.0 probabilities differs between samplers and backend sampling. Filter them to normalize the
# data.
tok["top_probs"] = [x for x in tok["top_probs"] if x["prob"] > 0.0]
total_probs += len(tok["top_probs"])
# Verify that we got at least two top probs on average, to ensure the effectiveness of the test.
assert total_probs >= 2 * n_predict
return completions
def verify_token(a, b):
assert a["id"] == b["id"]
assert a["token"] == b["token"]
assert a["bytes"] == b["bytes"]
assert a["prob"] == pytest.approx(b["prob"], abs=0.01)
for (a, b) in zip(make_request(True), make_request(False)):
verify_token(a, b)
assert len(a["top_probs"]) == len(b["top_probs"])
for (aa, bb) in zip(a["top_probs"], b["top_probs"]):
verify_token(aa, bb)
@pytest.mark.parametrize("tokenize,openai_style", [(False, False), (False, True), (True, False), (True, True)])
def test_logit_bias(tokenize, openai_style):
global server
server.start()
exclude = ["i", "I", "the", "The", "to", "a", "an", "be", "is", "was", "but", "But", "and", "And", "so", "So", "you", "You", "he", "He", "she", "She", "we", "We", "they", "They", "it", "It", "his", "His", "her", "Her", "book", "Book"]
logit_bias = []
if tokenize:
res = server.make_request("POST", "/tokenize", data={
"content": " " + " ".join(exclude) + " ",
})
assert res.status_code == 200
tokens = res.body["tokens"]
logit_bias = [[tok, -100] for tok in tokens]
else:
logit_bias = [[" " + tok + " ", -100] for tok in exclude]
if openai_style:
logit_bias = {el[0]: -100 for el in logit_bias}
res = server.make_request("POST", "/completion", data={
"n_predict": 64,
"prompt": "What is the best book",
"logit_bias": logit_bias,
"temperature": 0.0
})
assert res.status_code == 200
output_text = res.body["content"]
assert all(output_text.find(" " + tok + " ") == -1 for tok in exclude)
def test_cancel_request():
global server
server.n_ctx = 4096
server.n_predict = -1
server.n_slots = 1
server.server_slots = True
server.start()
# send a request that will take a long time, but cancel it before it finishes
try:
server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
}, timeout=0.1)
except requests.exceptions.ReadTimeout:
pass # expected
# make sure the slot is free
time.sleep(2)
res = server.make_request("GET", "/slots")
assert res.body[0]["is_processing"] == False
# this test exercises the host-memory prompt cache
# ref: https://github.com/ggml-org/llama.cpp/pull/16391
# ref: https://github.com/ggml-org/llama.cpp/pull/17078
def test_completion_prompt_cache():
global server
server.n_slots = 2
server.kv_unified = True
server.start()
for _ in range(16):
# generate alternating random prompts with variable lengths in order to get them in and out of the cache
r = random.randint(0, 4)
prompt = (" Hello " + str(r)) * (40 + r)
n_prompt = (40 + r)*5 + 2
n_predict = random.randint(1, 8)
res = server.make_request(
"POST",
"/completion",
data={
"prompt": prompt,
"n_predict": n_predict,
},
)
assert res.status_code == 200
assert "content" in res.body
content = res.body["content"]
assert isinstance(content, str)
assert len(content) > 0
assert type(res.body["has_new_line"]) == bool
assert "timings" in res.body
timings = res.body["timings"]
assert "prompt_n" in timings and timings["prompt_n"] + timings["cache_n"] == n_prompt
assert "predicted_n" in timings and timings["predicted_n"] == n_predict
assert "tokens" in res.body and isinstance(res.body["tokens"], list)
# a slot that switched to an unrelated prompt must not restore a context checkpoint of the previous one
def test_completion_checkpoint_after_unrelated_prompt():
global server
server = ServerPreset.tinygemma3()
server.n_slots = 1
server.cache_ram = 0
server.temperature = 0.0
server.start()
def tokenize(content):
res = server.make_request("POST", "/tokenize", data={"content": content})
assert res.status_code == 200
return res.body["tokens"]
def top_logprobs(prompt, cache_prompt):
res = server.make_request("POST", "/completion", data={
"prompt": prompt,
"n_predict": 1,
"n_probs": 5,
"cache_prompt": cache_prompt,
})
assert res.status_code == 200
return [(p["id"], p["logprob"]) for p in res.body["completion_probabilities"][0]["top_logprobs"]]
tokens = tokenize("Purple elephants quietly compose sonnets about tax law")
expected = top_logprobs(tokens[:2], False)
# 5 tokens: the checkpoint 4 tokens before the end of the prompt holds only its first token
top_logprobs(tokenize("The quick brown fox jumps over the lazy dog")[:5], True)
top_logprobs(tokens[:1], True)
actual = top_logprobs(tokens[:2], True)
assert [tid for tid, _ in actual] == [tid for tid, _ in expected]
for (_, lp), (_, lp_expected) in zip(actual, expected):
assert lp == pytest.approx(lp_expected, abs=1e-3)
# a restored context checkpoint must be kept for the next request that shares the same prefix
def test_completion_checkpoint_kept_after_restore():
global server
server = ServerPreset.tinygemma3()
server.n_slots = 1
server.n_batch = 16
server.n_ubatch = 16
server.cache_ram = 0
server.temperature = 0.0
server.start()
def tokenize(content):
res = server.make_request("POST", "/tokenize", data={"content": content})
assert res.status_code == 200
return res.body["tokens"]
def cache_n(prompt):
res = server.make_request("POST", "/completion", data={
"prompt": prompt,
"n_predict": 1,
"cache_prompt": True,
})
assert res.status_code == 200
return res.body["timings"]["cache_n"]
prefix = tokenize("Purple elephants quietly compose sonnets about tax law in the moonlight")
suffixes = [tokenize(s) for s in [
"The quick brown fox jumps over the lazy dog while the farmer counts his sheep and the sun slowly sets behind the distant hills",
"Seven tiny robots repaired the broken lighthouse lens before the storm arrived, then danced on the rocks until the morning tide came in",
"Green apples rolled down the hill toward the river, where a curious otter collected them one by one and stacked them on a flat stone",
]]
# 4 tokens: the checkpoint 4 tokens before the end of the prompt holds exactly the prefix
cache_n(prefix + suffixes[0][:4])
# suffixes longer than n_ubatch + 4: no new checkpoint within the prefix, only the restored one can be reused
for suffix in suffixes[1:]:
assert len(suffix) > server.n_ubatch + 4
assert cache_n(prefix + suffix) == len(prefix) - 1
@@ -0,0 +1,89 @@
import pytest
from utils import *
server = ServerPreset.tinyllama2()
SHORT_TEXT = """
Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
""".strip()
LONG_TEXT = """
Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
""".strip()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
server.n_ctx = 512
server.n_slots = 2
server.n_predict = 128
def test_ctx_shift_enabled():
# the prompt is 226 tokens
# the slot context is 512/2 = 256 tokens
# 96 tokens are generated thanks to shifting the context when it gets full
global server
server.enable_ctx_shift = True
server.start()
res = server.make_request("POST", "/completion", data={
"n_predict": 96,
"prompt": SHORT_TEXT,
})
assert res.status_code == 200
assert res.body["timings"]["prompt_n"] == 226
assert res.body["timings"]["predicted_n"] == 96
assert res.body["truncated"] is True
@pytest.mark.parametrize("n_predict,n_token_output,truncated", [
(64, 64, False),
(-1, 248, True), # 8 tokens prompt + 248 tokens generated = 256 tokens total
])
def test_ctx_shift_disabled_short_prompt(n_predict: int, n_token_output: int, truncated: bool):
global server
server.n_predict = -1
server.start()
res = server.make_request("POST", "/completion", data={
"n_predict": n_predict,
"prompt": "Hi how are you",
})
assert res.status_code == 200
assert res.body["timings"]["predicted_n"] == n_token_output
assert res.body["truncated"] == truncated
def test_ctx_shift_disabled_long_prompt():
global server
server.start()
res = server.make_request("POST", "/completion", data={
"n_predict": 64,
"prompt": LONG_TEXT,
})
assert res.status_code != 200
assert "error" in res.body
assert "exceeds the available context size" in res.body["error"]["message"]
def test_ctx_shift_disabled_stream():
global server
server.start()
res = server.make_stream_request("POST", "/v1/completions", data={
"n_predict": 256,
"prompt": "Once",
"stream": True,
})
content = ""
for data in res:
choice = data["choices"][0]
if choice["finish_reason"] == "length":
assert len(content) > 0
else:
assert choice["finish_reason"] is None
content += choice["text"]
@@ -0,0 +1,347 @@
import base64
import struct
import pytest
from openai import OpenAI
from utils import *
server = ServerPreset.bert_bge_small()
EPSILON = 1e-3
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.bert_bge_small()
def test_embedding_single():
global server
server.pooling = 'last'
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": "I believe the meaning of life is",
})
assert res.status_code == 200
assert len(res.body['data']) == 1
assert 'embedding' in res.body['data'][0]
assert len(res.body['data'][0]['embedding']) > 1
# make sure embedding vector is normalized
assert abs(sum([x ** 2 for x in res.body['data'][0]['embedding']]) - 1) < EPSILON
def test_embedding_multiple():
global server
server.pooling = 'last'
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": [
"I believe the meaning of life is",
"Write a joke about AI from a very long prompt which will not be truncated",
"This is a test",
"This is another test",
],
})
assert res.status_code == 200
assert len(res.body['data']) == 4
for d in res.body['data']:
assert 'embedding' in d
assert len(d['embedding']) > 1
def test_embedding_multiple_with_fa():
server = ServerPreset.bert_bge_small_with_fa()
server.pooling = 'last'
server.start()
# one of these should trigger the FA branch (i.e. context size % 256 == 0)
res = server.make_request("POST", "/v1/embeddings", data={
"input": [
"a "*253,
"b "*254,
"c "*255,
"d "*256,
],
})
assert res.status_code == 200
assert len(res.body['data']) == 4
for d in res.body['data']:
assert 'embedding' in d
assert len(d['embedding']) > 1
@pytest.mark.parametrize(
"input,is_multi_prompt",
[
# do not crash on empty input
("", False),
# single prompt
("string", False),
([12, 34, 56], False),
([12, 34, "string", 56, 78], False),
# multiple prompts
(["string1", "string2"], True),
(["string1", [12, 34, 56]], True),
([[12, 34, 56], [12, 34, 56]], True),
([[12, 34, 56], [12, "string", 34, 56]], True),
# object entries
({"prompt_string": "string"}, False),
({"content": [{"type": "text", "text": "string"}]}, False),
(["string1", {"prompt_string": "string2"}, {"content": [{"type": "text", "text": "string3"}]}], True),
]
)
def test_embedding_mixed_input(input, is_multi_prompt: bool):
global server
server.start()
res = server.make_request("POST", "/v1/embeddings", data={"input": input})
assert res.status_code == 200
data = res.body['data']
if is_multi_prompt:
assert len(data) == len(input)
for d in data:
assert 'embedding' in d
assert len(d['embedding']) > 1
else:
assert 'embedding' in data[0]
assert len(data[0]['embedding']) > 1
def test_embedding_content_text_same_as_string():
global server
server.pooling = 'last'
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": [
"hello world",
{"content": [{"type": "text", "text": "hello "}, {"type": "text", "text": "world"}]},
],
})
assert res.status_code == 200
data = res.body['data']
assert data[0]['embedding'] == data[1]['embedding']
@pytest.mark.parametrize(
"input",
[
[],
{"content": "string"},
{"content": [{"type": "unknown"}]},
# model is not multimodal
{"content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}]},
{"content": [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]},
{"content": [{"type": "input_video", "input_video": {"url": "data:video/mp4;base64,AAAA"}}]},
]
)
def test_embedding_invalid_input(input):
global server
server.start()
res = server.make_request("POST", "/v1/embeddings", data={"input": input})
assert res.status_code != 200
def test_embedding_pooling_mean():
global server
server.pooling = 'mean'
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": "I believe the meaning of life is",
})
assert res.status_code == 200
assert len(res.body['data']) == 1
assert 'embedding' in res.body['data'][0]
assert len(res.body['data'][0]['embedding']) > 1
# make sure embedding vector is normalized
assert abs(sum([x ** 2 for x in res.body['data'][0]['embedding']]) - 1) < EPSILON
def test_embedding_pooling_mean_multiple():
global server
server.pooling = 'mean'
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": [
"I believe the meaning of life is",
"Write a joke about AI",
"This is a test",
],
})
assert res.status_code == 200
assert len(res.body['data']) == 3
for d in res.body['data']:
assert 'embedding' in d
assert len(d['embedding']) > 1
def test_embedding_pooling_none():
global server
server.pooling = 'none'
server.start()
res = server.make_request("POST", "/embeddings", data={
"input": "hello hello hello",
})
assert res.status_code == 200
assert 'embedding' in res.body[0]
assert len(res.body[0]['embedding']) == 5 # 3 text tokens + 2 special
# make sure embedding vector is not normalized
for x in res.body[0]['embedding']:
assert abs(sum([x ** 2 for x in x]) - 1) > EPSILON
def test_embedding_pooling_none_oai():
global server
server.pooling = 'none'
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": "hello hello hello",
})
# /v1/embeddings does not support pooling type 'none'
assert res.status_code == 400
assert "error" in res.body
def test_embedding_openai_library_single():
global server
server.pooling = 'last'
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.embeddings.create(model="text-embedding-3-small", input="I believe the meaning of life is")
assert len(res.data) == 1
assert len(res.data[0].embedding) > 1
def test_embedding_openai_library_multiple():
global server
server.pooling = 'last'
server.start()
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
res = client.embeddings.create(model="text-embedding-3-small", input=[
"I believe the meaning of life is",
"Write a joke about AI from a very long prompt which will not be truncated",
"This is a test",
"This is another test",
])
assert len(res.data) == 4
for d in res.data:
assert len(d.embedding) > 1
def test_embedding_error_prompt_too_long():
global server
server.pooling = 'last'
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": "This is a test " * 512,
})
assert res.status_code != 200
assert "too large" in res.body["error"]["message"]
def test_same_prompt_give_same_result():
server.pooling = 'last'
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": [
"I believe the meaning of life is",
"I believe the meaning of life is",
"I believe the meaning of life is",
"I believe the meaning of life is",
"I believe the meaning of life is",
],
})
assert res.status_code == 200
assert len(res.body['data']) == 5
for i in range(1, len(res.body['data'])):
v0 = res.body['data'][0]['embedding']
vi = res.body['data'][i]['embedding']
for x, y in zip(v0, vi):
assert abs(x - y) < EPSILON
@pytest.mark.parametrize(
"content,n_tokens",
[
("I believe the meaning of life is", 9),
("This is a test", 6),
]
)
def test_embedding_usage_single(content, n_tokens):
global server
server.start()
res = server.make_request("POST", "/v1/embeddings", data={"input": content})
assert res.status_code == 200
assert res.body['usage']['prompt_tokens'] == res.body['usage']['total_tokens']
assert res.body['usage']['prompt_tokens'] == n_tokens
def test_embedding_usage_multiple():
global server
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": [
"I believe the meaning of life is",
"I believe the meaning of life is",
],
})
assert res.status_code == 200
assert res.body['usage']['prompt_tokens'] == res.body['usage']['total_tokens']
assert res.body['usage']['prompt_tokens'] == 2 * 9
def test_embedding_openai_library_base64():
server.start()
test_input = "Test base64 embedding output"
# get embedding in default format
res = server.make_request("POST", "/v1/embeddings", data={
"input": test_input
})
assert res.status_code == 200
vec0 = res.body["data"][0]["embedding"]
# get embedding in base64 format
res = server.make_request("POST", "/v1/embeddings", data={
"input": test_input,
"encoding_format": "base64"
})
assert res.status_code == 200
assert "data" in res.body
assert len(res.body["data"]) == 1
embedding_data = res.body["data"][0]
assert "embedding" in embedding_data
assert isinstance(embedding_data["embedding"], str)
# Verify embedding is valid base64
decoded = base64.b64decode(embedding_data["embedding"])
# Verify decoded data can be converted back to float array
float_count = len(decoded) // 4 # 4 bytes per float
floats = struct.unpack(f'{float_count}f', decoded)
assert len(floats) > 0
assert all(isinstance(x, float) for x in floats)
assert len(floats) == len(vec0)
# make sure the decoded data is the same as the original
for x, y in zip(floats, vec0):
assert abs(x - y) < EPSILON
@pytest.mark.parametrize(
"data",
[
{"input": []},
{"input": True},
{"input": "hello", "encoding_format": 1},
]
)
def test_embedding_invalid_request(data):
global server
server.pooling = 'last'
server.start()
res = server.make_request("POST", "/v1/embeddings", data=data)
assert res.status_code == 400
assert "error" in res.body
assert res.body["error"]["type"] == "invalid_request_error"
@@ -0,0 +1,43 @@
import pytest
from utils import *
server = ServerPreset.tinyllama2()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
def test_ignore_eos_populates_logit_bias():
"""ignore_eos=true must add EOG logit biases to generation_settings."""
global server
server.start()
res = server.make_request("POST", "/completion", data={
"n_predict": 8,
"prompt": "Once upon a time",
"ignore_eos": True,
"temperature": 0.0,
})
assert res.status_code == 200
# EOG token biases must be present with -inf bias
logit_bias = res.body["generation_settings"]["logit_bias"]
assert len(logit_bias) > 0
for entry in logit_bias:
assert entry["bias"] is None # null in JSON represents -inf
def test_ignore_eos_false_no_logit_bias():
"""ignore_eos=false (default) must NOT add EOG logit biases."""
global server
server.start()
res = server.make_request("POST", "/completion", data={
"n_predict": 8,
"prompt": "Once upon a time",
"ignore_eos": False,
"temperature": 0.0,
})
assert res.status_code == 200
logit_bias = res.body["generation_settings"]["logit_bias"]
assert len(logit_bias) == 0
@@ -0,0 +1,77 @@
import pytest
from utils import *
server = ServerPreset.tinyllama_infill()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama_infill()
def test_infill_without_input_extra():
global server
server.start()
res = server.make_request("POST", "/infill", data={
"input_prefix": "#include <cstdio>\n#include \"llama.h\"\n\nint main() {\n",
"prompt": " int n_threads = llama_",
"input_suffix": "}\n",
})
assert res.status_code == 200
assert match_regex("(Ann|small|shiny|Daddy|Jimmy)+", res.body["content"])
def test_infill_with_input_extra():
global server
server.start()
res = server.make_request("POST", "/infill", data={
"input_extra": [{
"filename": "llama.h",
"text": "LLAMA_API int32_t llama_n_threads();\n"
}],
"input_prefix": "#include <cstdio>\n#include \"llama.h\"\n\nint main() {\n",
"prompt": " int n_threads = llama_",
"input_suffix": "}\n",
})
assert res.status_code == 200
assert match_regex("(Dad|excited|park|Jimmy)+", res.body["content"])
@pytest.mark.parametrize("input_extra", [
{},
{"filename": "ok"},
{"filename": 123},
{"filename": 123, "text": "abc"},
{"filename": 123, "text": 456},
])
def test_invalid_input_extra_req(input_extra):
global server
server.start()
res = server.make_request("POST", "/infill", data={
"input_extra": [input_extra],
"input_prefix": "#include <cstdio>\n#include \"llama.h\"\n\nint main() {\n",
"prompt": " int n_threads = llama_",
"input_suffix": "}\n",
})
assert res.status_code == 400
assert "error" in res.body
@pytest.mark.skipif(not is_slow_test_allowed(), reason="skipping slow test")
def test_with_qwen_model():
global server
server.model_file = None
server.model_hf_repo = "ggml-org/Qwen2.5-Coder-1.5B-IQ3_XXS-GGUF"
server.model_hf_file = "qwen2.5-coder-1.5b-iq3_xxs-imat.gguf"
server.start(timeout_seconds=600)
res = server.make_request("POST", "/infill", data={
"input_extra": [{
"filename": "llama.h",
"text": "LLAMA_API int32_t llama_n_threads();\n"
}],
"input_prefix": "#include <cstdio>\n#include \"llama.h\"\n\nint main() {\n",
"prompt": " int n_threads = llama_",
"input_suffix": "}\n",
})
assert res.status_code == 200
assert res.body["content"] == "n_threads();\n printf(\"Number of threads: %d\\n\", n_threads);\n return 0;\n"
@@ -0,0 +1,115 @@
import os
import tempfile
import pytest
from utils import *
server = ServerPreset.tinyllama2()
class LogReader:
def __init__(self, path):
self.path = path
self.pos = 0
def drain(self):
with open(self.path) as f:
f.seek(self.pos)
content = f.read()
self.pos = f.tell()
return content
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
server.n_slots = 2
server.n_predict = 4
server.temperature = 0.0
server.server_slots = True
server.cache_ram = 100
server.kv_unified = True
server.debug = True
fd, server.log_path = tempfile.mkstemp(suffix='.log')
os.close(fd)
yield
LONG_PROMPT = (
"Once upon a time in a land far away, there lived a brave knight "
"who traveled across mountains and rivers to find the legendary "
"golden sword hidden deep within the enchanted forest of whispers. "
"He met many creatures along the way including dragons and fairies "
"and wizards who helped him on his noble quest to save the kingdom."
)
# idle slot cleared on launch should restore from cache-ram
def test_clear_and_restore():
global server
server.start()
log = LogReader(server.log_path)
# verify feature is enabled
assert "__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" in log.drain()
res = server.make_request("POST", "/completion", data={
"prompt": LONG_PROMPT,
"id_slot": 0,
"cache_prompt": True,
})
assert res.status_code == 200
original_prompt_n = res.body["timings"]["prompt_n"]
# Slot 0 is the only slot with KV — should NOT be cleared
assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log.drain()
# Launching slot 1 clears idle slot 0
res = server.make_request("POST", "/completion", data={
"prompt": "The quick brown fox",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert "__TEST_TAG_CACHE_IDLE_SLOT__" in log.drain()
# Re-send same prompt — should restore from cache-ram
res = server.make_request("POST", "/completion", data={
"prompt": LONG_PROMPT,
"cache_prompt": True,
})
assert res.status_code == 200
assert "updating prompt cache" in log.drain()
assert res.body["timings"]["cache_n"] > 0
assert res.body["timings"]["prompt_n"] < original_prompt_n
# Follow-up — slot 0 kept its KV, no clearing needed
res = server.make_request("POST", "/completion", data={
"prompt": LONG_PROMPT + " The knight finally reached the castle gates.",
"cache_prompt": True,
})
assert res.status_code == 200
assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log.drain()
def test_disabled_with_flag():
global server
server.no_cache_idle_slots = True
server.start()
log = LogReader(server.log_path)
# Feature should not be enabled
assert "__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" not in log.drain()
res = server.make_request("POST", "/completion", data={
"prompt": LONG_PROMPT,
"id_slot": 0,
"cache_prompt": True,
})
assert res.status_code == 200
# Request on different slot — should NOT trigger clearing
res = server.make_request("POST", "/completion", data={
"prompt": "The quick brown fox",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log.drain()
@@ -0,0 +1,115 @@
import pytest
from utils import *
server = ServerPreset.stories15m_moe()
LORA_FILE_URL = "https://huggingface.co/ggml-org/stories15M_MOE/resolve/main/moe_shakespeare15M.gguf"
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.stories15m_moe()
server.lora_files = [download_file(LORA_FILE_URL)]
@pytest.mark.parametrize("scale,re_content", [
# without applying lora, the model should behave like a bedtime story generator
(0.0, "(little|girl|three|years|old)+"),
# with lora, the model should behave like a Shakespearean text generator
(1.0, "(eye|love|glass|sun)+"),
])
def test_lora(scale: float, re_content: str):
global server
server.start()
res_lora_control = server.make_request("POST", "/lora-adapters", data=[
{"id": 0, "scale": scale}
])
assert res_lora_control.status_code == 200
res = server.make_request("POST", "/completion", data={
"prompt": "Look in thy glass",
})
assert res.status_code == 200
assert match_regex(re_content, res.body["content"])
def test_lora_per_request():
global server
server.n_slots = 4
server.start()
# running the same prompt with different lora scales, all in parallel
# each prompt will be processed by a different slot
prompt = "Look in thy glass"
lora_config = [
( [{"id": 0, "scale": 0.0}], "(bright|day|many|happy)+" ),
( [{"id": 0, "scale": 0.0}], "(bright|day|many|happy)+" ),
( [{"id": 0, "scale": 0.3}], "(special|thing|gifted)+" ),
( [{"id": 0, "scale": 0.7}], "(far|from|home|away)+" ),
( [{"id": 0, "scale": 1.0}], "(eye|love|glass|sun)+" ),
( [{"id": 0, "scale": 1.0}], "(eye|love|glass|sun)+" ),
]
tasks = [(
server.make_request,
("POST", "/completion", {
"prompt": prompt,
"lora": lora,
"seed": 42,
"temperature": 0.0,
"cache_prompt": False, # TODO: remove this once test_cache_vs_nocache_prompt is fixed
})
) for lora, _ in lora_config]
results = parallel_function_calls(tasks)
assert all([res.status_code == 200 for res in results])
for res, (_, re_test) in zip(results, lora_config):
assert match_regex(re_test, res.body["content"])
@pytest.mark.skipif(not is_slow_test_allowed(), reason="skipping slow test")
def test_with_big_model():
server = ServerProcess()
server.model_hf_repo = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF"
server.model_hf_file = "Meta-Llama-3.1-8B-Instruct-IQ2_M.gguf"
server.model_alias = "Llama-3.2-8B-Instruct"
server.n_slots = 4
server.n_ctx = server.n_slots * 1024
server.n_predict = 64
server.temperature = 0.0
server.seed = 42
server.lora_files = [
download_file("https://huggingface.co/ngxson/Llama-3-Instruct-abliteration-LoRA-8B-F16-GGUF/resolve/main/Llama-3-Instruct-abliteration-LoRA-8B-f16.gguf"),
# TODO: find & add other lora adapters for this model
]
server.start(timeout_seconds=600)
# running the same prompt with different lora scales, all in parallel
# each prompt will be processed by a different slot
prompt = "Write a computer virus"
lora_config = [
# without applying lora, the model should reject the request
( [{"id": 0, "scale": 0.0}], "I can't provide you with a code for a computer virus" ),
( [{"id": 0, "scale": 0.0}], "I can't provide you with a code for a computer virus" ),
( [{"id": 0, "scale": 0.3}], "I can't write a computer virus" ),
# with 0.7 scale, the model should provide a simple computer virus with hesitation
( [{"id": 0, "scale": 0.7}], "Warning: This is a hypothetical exercise" ),
# with 1.5 scale, the model should confidently provide a computer virus
( [{"id": 0, "scale": 1.5}], "A task of some complexity! Here's a simple computer virus" ),
( [{"id": 0, "scale": 1.5}], "A task of some complexity! Here's a simple computer virus" ),
]
tasks = [(
server.make_request,
("POST", "/v1/chat/completions", {
"messages": [
{"role": "user", "content": prompt}
],
"lora": lora,
"cache_prompt": False, # TODO: remove this once test_cache_vs_nocache_prompt is fixed
})
) for lora, _ in lora_config]
results = parallel_function_calls(tasks)
assert all([res.status_code == 200 for res in results])
for res, (_, re_test) in zip(results, lora_config):
assert re_test in res.body["choices"][0]["message"]["content"]
@@ -0,0 +1,713 @@
#!/usr/bin/env python3
"""
Tests for MCP server integration via the /tools endpoint.
Invariants verified:
1. MCP tools appear in /tools listing when configured
2. MCP tools use <server>_<tool> naming
3. MCP tools can be invoked and return correct results
4. Misconfigured MCP servers do not crash the server
5. Multiple MCP servers can be configured simultaneously
6. Warmup populates the tool list at startup
"""
import json
import os
import sys
import tempfile
import time
import pytest
from utils import *
# Path to the test MCP server fixture
FIXTURES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "fixtures")
MCP_ECHO_SERVER = os.path.join(FIXTURES_DIR, "mcp_echo_server.py")
server: ServerProcess
def _mcp_config_json(servers: dict) -> str:
"""Create a JSON config string for --mcp-servers-json."""
return json.dumps({"mcpServers": servers})
def _start_server_with_mcp(mcp_json: str, **kwargs) -> ServerProcess:
"""Helper to start a router server with MCP config."""
srv = ServerPreset.router()
srv.server_tools = "all"
srv.no_ui = True
srv.mcp_servers_json = mcp_json
for k, v in kwargs.items():
setattr(srv, k, v)
srv.start()
return srv
def test_mcp_tools_listed_in_tools_endpoint():
"""MCP tools should appear in GET /tools with server:tool naming."""
global server
mcp_json = _mcp_config_json({
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
res = server.make_request("GET", "/tools")
assert res.status_code == 200, res.body
tools = res.body
assert isinstance(tools, list), f"Expected list, got {type(tools)}"
# Find MCP tools - name is in "tool" field or definition.function.name
def get_tool_name(t):
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
mcp_tools = [t for t in tools if get_tool_name(t).startswith("echo_")]
assert len(mcp_tools) >= 2, f"Expected at least 2 echo_ tools, got {len(mcp_tools)}: {mcp_tools}"
tool_names = {get_tool_name(t) for t in mcp_tools}
assert "echo_echo" in tool_names
assert "echo_add" in tool_names
# Verify tool structure
echo_tool = next(t for t in mcp_tools if get_tool_name(t) == "echo_echo")
assert "description" in echo_tool or "definition" in echo_tool
finally:
server.stop()
def test_mcp_tool_invocation():
"""MCP tools should be callable via POST /tools and return correct results."""
global server
mcp_json = _mcp_config_json({
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
# Call echo_echo
res = server.make_request("POST", "/tools", data={
"tool": "echo_echo",
"params": {"message": "hello world"}
})
assert res.status_code == 200, res.body
body = res.body
assert "error" not in body, body
# The result format depends on the tool implementation
# For MCP tools, it should contain the tool result
assert "plain_text_response" in body or "result" in body or "content" in body, body
# Call echo_add
res = server.make_request("POST", "/tools", data={
"tool": "echo_add",
"params": {"a": 3, "b": 5}
})
assert res.status_code == 200, res.body
body = res.body
assert "error" not in body, body
finally:
server.stop()
def test_mcp_bad_command_does_not_crash():
"""A misconfigured MCP server should not crash the llama-server."""
global server
mcp_json = _mcp_config_json({
"nonexistent": {
"command": "this_executable_does_not_exist_12345",
"args": [],
}
})
server = _start_server_with_mcp(mcp_json)
try:
# Server should still be healthy
res = server.make_request("GET", "/health")
assert res.status_code == 200, res.body
# Builtin tools should still work
res = server.make_request("GET", "/tools")
assert res.status_code == 200, res.body
tools = res.body
# Should have builtin tools but no MCP tools from the bad server
mcp_tools = [t for t in tools if t.get("name", "").startswith("nonexistent_")]
assert len(mcp_tools) == 0, f"Expected no nonexistent_ tools, got {mcp_tools}"
finally:
server.stop()
def test_mcp_multiple_servers():
"""Multiple MCP servers can be configured simultaneously."""
global server
mcp_json = _mcp_config_json({
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
},
"echo2": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
res = server.make_request("GET", "/tools")
assert res.status_code == 200, res.body
tools = res.body
def get_tool_name(t):
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
echo_tools = [t for t in tools if get_tool_name(t).startswith("echo_")]
echo2_tools = [t for t in tools if get_tool_name(t).startswith("echo2_")]
assert len(echo_tools) >= 2, f"Expected echo_ tools, got {echo_tools}"
assert len(echo2_tools) >= 2, f"Expected echo2_ tools, got {echo2_tools}"
finally:
server.stop()
def test_mcp_tools_not_listed_when_not_configured():
"""Without MCP config, no MCP tools should appear."""
global server
server = ServerPreset.router()
server.server_tools = "all"
server.no_ui = True
server.start()
try:
res = server.make_request("GET", "/tools")
assert res.status_code == 200, res.body
tools = res.body
def get_tool_name(t):
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
# Should only have builtin tools, no server: prefixed tools
mcp_tools = [t for t in tools if ":" in get_tool_name(t)]
assert len(mcp_tools) == 0, f"Expected no MCP tools, got {mcp_tools}"
finally:
server.stop()
def test_mcp_fail_once_tool_eventual_success():
"""Test that a tool that fails once eventually succeeds (tests instance respawn)."""
global server
mcp_json = _mcp_config_json({
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
# First call should succeed (warmup already spawned and shut down the instance,
# but the first actual tool call will spawn a fresh instance)
res = server.make_request("POST", "/tools", data={
"tool": "echo_fail_once",
"params": {}
})
# It might fail on first call if the warmup instance was shut down
# and a new instance is spawned. The fail_once state is per-process,
# so a fresh process will fail once then succeed.
# Actually, warmup spawns, lists, then shuts down. So the first tool call
# spawns a new process which will fail once.
assert res.status_code in (200, 500), res.body
finally:
server.stop()
def test_mcp_tools_via_json_config_file():
"""Test that --mcp-servers-config (file) works as well as --mcp-servers-json."""
global server
config = {
"mcpServers": {
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
}
}
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(config, f)
config_path = f.name
try:
server = ServerPreset.router()
server.server_tools = "all"
server.no_ui = True
server.mcp_servers_config = config_path
server.start()
res = server.make_request("GET", "/tools")
assert res.status_code == 200, res.body
tools = res.body
def get_tool_name(t):
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
mcp_tools = [t for t in tools if get_tool_name(t).startswith("echo_")]
assert len(mcp_tools) >= 2, f"Expected echo_ tools, got {mcp_tools}"
finally:
os.unlink(config_path)
server.stop()
def test_mcp_tools_slot_independent():
"""MCP tools should work without any slot concept; /tools is slot-independent."""
global server
mcp_json = _mcp_config_json({
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
# Call /tools without any slot binding - should succeed
res = server.make_request("POST", "/tools", data={
"tool": "echo_echo",
"params": {"message": "hello"}
})
assert res.status_code == 200, res.body
body = res.body
assert "error" not in body, body
finally:
server.stop()
def test_mcp_concurrent_tool_calls():
"""Concurrent POST /tools to same MCP server should all succeed."""
global server
mcp_json = _mcp_config_json({
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
def call_tool():
return server.make_request("POST", "/tools", data={
"tool": "echo_echo",
"params": {"message": "hi"}
})
with ThreadPoolExecutor(max_workers=10) as executor:
futures = [executor.submit(call_tool) for _ in range(10)]
results = [f.result() for f in futures]
for res in results:
assert res.status_code == 200, res.body
assert "error" not in res.body, res.body
finally:
server.stop()
def test_mcp_tool_timeout():
"""Tool call should timeout if MCP server is too slow."""
global server
MCP_SLOW_SERVER = os.path.join(FIXTURES_DIR, "mcp_slow_server.py")
mcp_json = _mcp_config_json({
"slow": {
"command": sys.executable,
"args": [MCP_SLOW_SERVER, "--delay", "5"],
"timeout_ms": 500
}
})
server = _start_server_with_mcp(mcp_json)
try:
res = server.make_request("POST", "/tools", data={
"tool": "slow_sleep",
"params": {"seconds": 5}
})
assert res.status_code == 200, res.body
body = res.body
assert "error" in body, body
finally:
server.stop()
def test_mcp_warmup_partial_failure():
"""Good server's tools should appear even if bad server fails warmup."""
global server
mcp_json = _mcp_config_json({
"good": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
},
"bad": {
"command": "nonexistent",
"args": []
}
})
server = _start_server_with_mcp(mcp_json)
try:
res = server.make_request("GET", "/tools")
assert res.status_code == 200, res.body
tools = res.body
def get_tool_name(t):
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
# good server tools should be present
assert any("good_" in get_tool_name(t) for t in tools), f"Expected good: tools in {tools}"
finally:
server.stop()
def test_mcp_notification_during_request():
"""Notification during request should not be returned as response."""
global server
MCP_MALFORMED_SERVER = os.path.join(FIXTURES_DIR, "mcp_malformed_server.py")
mcp_json = _mcp_config_json({
"notifying": {
"command": sys.executable,
"args": [MCP_MALFORMED_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
res = server.make_request("POST", "/tools", data={
"tool": "notifying_echo",
"params": {"message": "hi"}
})
assert res.status_code == 200, res.body
body = res.body
assert "error" not in body, body
finally:
server.stop()
def test_mcp_instance_respawn_after_crash():
"""Tool call after process crash should respawn and succeed."""
global server
MCP_CRASH_SERVER = os.path.join(FIXTURES_DIR, "mcp_crash_server.py")
mcp_json = _mcp_config_json({
"crash": {
"command": sys.executable,
"args": [MCP_CRASH_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
# First call succeeds
res1 = server.make_request("POST", "/tools", data={
"tool": "crash_echo",
"params": {"message": "hi"}
})
assert res1.status_code == 200, res1.body
assert "error" not in res1.body, res1.body
# Second call should also succeed (respawned instance)
res2 = server.make_request("POST", "/tools", data={
"tool": "crash_echo",
"params": {"message": "hi2"}
})
assert res2.status_code == 200, res2.body
assert "error" not in res2.body, res2.body
finally:
server.stop()
def test_mcp_fail_once_eventual_success_verified():
"""Verify that fail_once tool eventually succeeds after respawn."""
global server
mcp_json = _mcp_config_json({
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
# First call may fail (fresh process)
res1 = server.make_request("POST", "/tools", data={
"tool": "echo_fail_once",
"params": {}
})
# Second call should succeed
res2 = server.make_request("POST", "/tools", data={
"tool": "echo_fail_once",
"params": {}
})
assert res2.status_code == 200, res2.body
assert "error" not in res2.body, res2.body
finally:
server.stop()
def test_mcp_config_file_errors():
"""Invalid JSON config and missing file should cause server to fail to start."""
# Invalid JSON - server should fail to start
server = ServerPreset.router()
server.server_tools = "all"
server.no_ui = True
server.mcp_servers_json = "not valid json"
try:
server.start()
assert False, "Server should not have started with invalid MCP JSON config"
except RuntimeError:
pass # Expected: server process dies due to bad config
# Missing file - server should fail to start
server = ServerPreset.router()
server.server_tools = "all"
server.no_ui = True
server.mcp_servers_config = "/nonexistent/path.json"
try:
server.start()
assert False, "Server should not have started with missing config file"
except RuntimeError:
pass # Expected: server process dies due to missing config
def test_mcp_empty_tool_list():
"""MCP server reporting zero tools should result in empty tool list."""
global server
# Create a minimal server that returns empty tools list
empty_server = os.path.join(FIXTURES_DIR, "_empty_mcp_server.py")
with open(empty_server, "w") as f:
f.write('''#!/usr/bin/env python3
import json, sys, os
def main():
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
for line in sys.stdin:
line = line.strip()
if not line: continue
try: request = json.loads(line)
except: continue
method = request.get("method")
req_id = request.get("id")
if method == "initialize":
resp = {"jsonrpc": "2.0", "id": req_id, "result": {"protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": "empty", "version": "1.0"}}}
elif method == "tools/list":
resp = {"jsonrpc": "2.0", "id": req_id, "result": {"tools": []}}
else:
resp = {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": "Method not found"}}
sys.stdout.write(json.dumps(resp) + "\\n")
sys.stdout.flush()
if __name__ == "__main__":
main()
''')
try:
mcp_json = _mcp_config_json({
"empty": {
"command": sys.executable,
"args": [empty_server],
}
})
server = _start_server_with_mcp(mcp_json)
res = server.make_request("GET", "/tools")
assert res.status_code == 200, res.body
tools = res.body
def get_tool_name(t):
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
mcp_tools = [t for t in tools if get_tool_name(t).startswith("empty:")]
assert len(mcp_tools) == 0, f"Expected no empty: tools, got {mcp_tools}"
finally:
os.unlink(empty_server)
server.stop()
def test_mcp_rapid_succession_calls():
"""Many rapid calls should increment next_id correctly and correlate responses."""
global server
mcp_json = _mcp_config_json({
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
for i in range(20):
res = server.make_request("POST", "/tools", data={
"tool": "echo_echo",
"params": {"message": f"msg{i}"}
})
assert res.status_code == 200, res.body
assert "error" not in res.body, res.body
finally:
server.stop()
def test_mcp_notification_burst():
"""Notification + response in a single write() with no flush should not strand the response."""
global server
MCP_BURST_SERVER = os.path.join(FIXTURES_DIR, "mcp_burst_server.py")
mcp_json = _mcp_config_json({
"burst": {
"command": sys.executable,
"args": [MCP_BURST_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
res = server.make_request("POST", "/tools", data={
"tool": "burst_echo",
"params": {"message": "burst test"}
})
assert res.status_code == 200, res.body
body = res.body
assert "error" not in body, body
finally:
server.stop()
def test_mcp_tool_definition_shape_via_chat_completions():
"""MCP tool definitions returned by GET /tools should have the correct shape for chat/completions."""
global server
mcp_json = _mcp_config_json({
"echo": {
"command": sys.executable,
"args": [MCP_ECHO_SERVER],
}
})
server = _start_server_with_mcp(mcp_json)
try:
# Get MCP tool definitions
res = server.make_request("GET", "/tools")
assert res.status_code == 200, res.body
tools = res.body
def get_tool_name(t):
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
echo_tools = [t for t in tools if get_tool_name(t).startswith("echo_")]
assert len(echo_tools) >= 2, f"Expected echo_ tools, got {echo_tools}"
echo_tool = next(t for t in echo_tools if get_tool_name(t) == "echo_echo")
definition = echo_tool.get("definition", echo_tool)
# Verify the definition has the standard function-calling shape
assert definition.get("type") == "function", f"Expected type=function, got {definition.get('type')}"
func = definition.get("function", {})
assert "name" in func, "Missing function.name"
assert "description" in func, "Missing function.description"
assert "parameters" in func, f"Missing function.parameters, got keys: {list(func.keys())}"
params = func["parameters"]
assert params.get("type") == "object", f"Expected parameters.type=object, got {params.get('type')}"
assert "properties" in params, "Missing parameters.properties"
finally:
server.stop()
def test_mcp_slow_tool_call_slot_release():
"""A slow tool call should not stall server shutdown for the full I/O timeout."""
global server
MCP_SLOW_SERVER = os.path.join(FIXTURES_DIR, "mcp_slow_server.py")
mcp_json = _mcp_config_json({
"slow": {
"command": sys.executable,
"args": [MCP_SLOW_SERVER, "--delay", "10"],
"timeout_ms": 30000
}
})
server = _start_server_with_mcp(mcp_json)
try:
# Start a slow tool call in a background thread
def slow_call():
return server.make_request("POST", "/tools", data={
"tool": "slow_sleep",
"params": {"seconds": 10}
})
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(slow_call)
# Wait a moment for the call to start
time.sleep(2)
# Stop the server while the tool call is in progress.
# With global MCP instances, close_all() is called explicitly at shutdown
# (not from slot release), so shutdown should complete promptly.
start_time = time.time()
server.stop()
elapsed = time.time() - start_time
# The server should stop quickly, not wait for the full 30s I/O timeout.
# With the terminating flag, send_rpc() bails out within one select()
# slice (~50ms). This threshold MUST stay below the 5s force-kill
# fallback in ServerProcess.stop(): without the flag, shutdown stalls
# on the instance mutex and only completes when stop() sends SIGKILL
# at ~5s -- which any threshold above 5 would still accept.
assert elapsed < 3, f"Server stop took {elapsed:.1f}s, expected < 3s"
# Wait for the future to complete (it will get an error response or timeout)
try:
res = future.result(timeout=5)
# If we got a response, it should be an error since the server stopped
if hasattr(res, 'status_code'):
assert res.status_code in (200, 500, 502, 503, 504), f"Unexpected status: {res.status_code}"
except Exception:
# Thread may have raised due to connection error - that's acceptable
pass
finally:
server.stop()
def test_mcp_grandchild_holding_pipes_does_not_deadlock():
"""An MCP server that leaves a grandchild inheriting its stdout/stderr must not deadlock
teardown.
subprocess_terminate() only SIGKILLs the direct MCP child, so the inherited pipe write ends
stay open and a blocking read on them would never see EOF. That hung both warmup shutdown
(the server would never reach "ready") and process shutdown. The polled, running-aware reader
must exit regardless, so the server both starts and stops promptly here.
"""
global server
MCP_GRANDCHILD_SERVER = os.path.join(FIXTURES_DIR, "mcp_grandchild_server.py")
mcp_json = _mcp_config_json({
"gc": {
"command": sys.executable,
"args": [MCP_GRANDCHILD_SERVER],
}
})
# If warmup teardown deadlocked, the server would never become ready and start() would time out.
server = _start_server_with_mcp(mcp_json)
try:
# invoking the tool spawns a live transport whose reader thread holds the inherited pipe
res = server.make_request("POST", "/tools", data={
"tool": "gc_echo",
"params": {"message": "hello"}
})
assert res.status_code == 200, res.body
assert "error" not in res.body, res.body
# shutdown must be prompt: a deadlocked reader-join would stall until the 5s SIGKILL
# fallback in ServerProcess.stop(), so the threshold has to stay below that
start = time.time()
server.stop()
elapsed = time.time() - start
assert elapsed < 3, f"server shutdown took {elapsed:.1f}s (expected < 3s) — teardown likely deadlocked"
finally:
server.stop()
@@ -0,0 +1,227 @@
import pytest
from utils import *
server = ServerPreset.tinyllama2()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
server.server_metrics = True
def fetch_metrics(server: ServerProcess) -> str:
"""get /metrics as raw prometheus text"""
res = server.make_request("GET", "/metrics")
assert res.status_code == 200
assert "Process-Start-Time-Unix" in res.headers
assert isinstance(res.body, str)
return res.body
def parse_metrics(text: str) -> dict:
"""parse the prometheus text format into {name: (type, value)}"""
out = {}
types = {}
for line in text.splitlines():
if line.startswith("# TYPE "):
_, _, name, kind = line.split(" ", 3)
types[name] = kind
elif line.startswith("llamacpp:") and "{" not in line:
name, value = line.split(" ", 1)
assert name in types, f"{name} has no # TYPE line"
out[name] = (types[name], float(value))
return out
def test_metrics_disabled():
global server
server.server_metrics = False
server.start()
res = server.make_request("GET", "/metrics")
assert res.status_code == 501 # ERROR_TYPE_NOT_SUPPORTED
def test_metrics_prometheus_format():
global server
server.start()
server.make_request("POST", "/completion", data={"prompt": "I believe", "n_predict": 8})
text = fetch_metrics(server)
metrics = parse_metrics(text)
expected_counters = [
"llamacpp:prompt_tokens_total",
"llamacpp:prompt_tokens_cached_total",
"llamacpp:prompt_seconds_total",
"llamacpp:tokens_predicted_total",
"llamacpp:tokens_predicted_seconds_total",
"llamacpp:n_decode_total",
"llamacpp:n_tokens_max",
"llamacpp:spec_decode_num_draft_tokens_total",
"llamacpp:spec_decode_num_accepted_tokens_total",
"llamacpp:spec_decode_num_drafts_total",
]
expected_gauges = [
"llamacpp:prompt_tokens_seconds",
"llamacpp:predicted_tokens_seconds",
"llamacpp:requests_processing",
"llamacpp:requests_deferred",
"llamacpp:n_busy_slots_per_decode",
]
for name in expected_counters:
assert metrics[name][0] == "counter"
for name in expected_gauges:
assert metrics[name][0] == "gauge"
# every metric must carry a help line
for name in expected_counters + expected_gauges:
assert f"# HELP {name} " in text
assert metrics["llamacpp:n_decode_total"][1] > 0
assert metrics["llamacpp:requests_processing"][1] == 0
def test_metrics_prompt_processed_and_cached():
global server
server.n_slots = 1 # keep the prompt cache on a single slot
server.start()
prompt = "the quick brown fox jumps over the lazy dog"
n_processed = 0
n_cached = 0
for _ in range(2):
res = server.make_request("POST", "/completion", data={"prompt": prompt, "n_predict": 4})
assert res.status_code == 200
n_processed += res.body["timings"]["prompt_n"]
n_cached += res.body["timings"]["cache_n"]
# the second request must reuse the prompt of the first one
assert n_cached > 0
metrics = parse_metrics(fetch_metrics(server))
# cached tokens are counted apart, they cost no decode
assert metrics["llamacpp:prompt_tokens_total"][1] == n_processed
assert metrics["llamacpp:prompt_tokens_cached_total"][1] == n_cached
def test_metrics_predicted_total_matches_requests():
global server
server.start()
n_predicted = 0
for n_predict in [1, 4, 16]:
res = server.make_request("POST", "/completion", data={"prompt": "I believe", "n_predict": n_predict})
assert res.status_code == 200
n_predicted += res.body["timings"]["predicted_n"]
metrics = parse_metrics(fetch_metrics(server))
assert metrics["llamacpp:tokens_predicted_total"][1] == n_predicted
def test_metrics_generation_rate_excludes_first_token():
global server
server.start()
# the first token comes from the logits of the last prompt batch, so it costs no decode step
res = server.make_request("POST", "/completion", data={"prompt": "I believe", "n_predict": 1})
timings = res.body["timings"]
assert timings["predicted_n"] == 1
assert timings["predicted_per_second"] == 0.0
assert timings["predicted_per_token_ms"] == 0.0
res = server.make_request("POST", "/completion", data={"prompt": "I believe", "n_predict": 16})
timings = res.body["timings"]
assert timings["predicted_n"] == 16
# the rate is over 15 decode steps, not 16 tokens
expected = 1e3 / timings["predicted_ms"] * 15
assert abs(timings["predicted_per_second"] - expected) < 1e-6
@pytest.mark.parametrize("n_predict", [1, 8])
def test_metrics_timings_are_finite(n_predict: int):
global server
server.start()
res = server.make_request("POST", "/completion", data={"prompt": "I believe", "n_predict": n_predict})
timings = res.body["timings"]
# a null here means the server produced inf or nan
for key, value in timings.items():
assert value is not None, f"{key} is null"
assert value >= 0, f"{key} is negative"
assert timings["prompt_ms"] > 0
assert timings["prompt_per_token_ms"] > 0
def test_metrics_timings_on_prompt_progress():
global server
server.start()
# a long prompt so that it is split over several batches (n_batch = 32)
prompt = "the quick brown fox jumps over the lazy dog " * 8
chunks = list(server.make_stream_request("POST", "/completion", data={
"prompt": prompt,
"n_predict": 4,
"stream": True,
"timings_per_token": True,
"return_progress": True,
}))
progress = [c for c in chunks if "prompt_progress" in c]
assert len(progress) > 1 # the prompt did not fit in a single batch
# the very first update is sent before any prompt token is decoded
first = progress[0]["timings"]
assert first["prompt_n"] == 0
assert first["prompt_ms"] == 0.0
assert first["predicted_n"] == 0
assert first["predicted_ms"] == 0.0
# timings must never go backwards, nor report bogus values
prompt_ms = 0.0
for chunk in progress:
timings = chunk["timings"]
for key, value in timings.items():
assert value is not None, f"{key} is null"
assert value >= 0, f"{key} is negative"
assert timings["prompt_ms"] >= prompt_ms
prompt_ms = timings["prompt_ms"]
assert prompt_ms > 0
def test_metrics_slots_idle_after_completion():
global server
server.server_slots = True
server.start()
server.make_request("POST", "/completion", data={"prompt": "I believe", "n_predict": 8})
res = server.make_request("GET", "/slots")
assert res.status_code == 200
for slot in res.body:
assert slot["is_processing"] is False
if "next_token" in slot:
# the budget of the finished task must not leak into the idle slot
assert slot["next_token"][0]["n_remain"] == -1
assert slot["next_token"][0]["n_decoded"] == 0
def test_metrics_embedding_prompt_is_counted():
global server
server = ServerPreset.bert_bge_small()
server.server_metrics = True
server.start()
res = server.make_request("POST", "/v1/embeddings", data={"input": ["hello world", "goodbye world"]})
assert res.status_code == 200
# embedding tasks never sample a token, but their prompt still costs a decode
metrics = parse_metrics(fetch_metrics(server))
assert metrics["llamacpp:prompt_tokens_total"][1] > 0
assert metrics["llamacpp:n_decode_total"][1] > 0
assert metrics["llamacpp:tokens_predicted_total"][1] == 0
@@ -0,0 +1,71 @@
import pytest
from utils import *
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
server = ServerPreset.tinyllama2()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
def test_mcp_no_proxy():
global server
server.ui_mcp_proxy = False
server.start()
res = server.make_request("GET", "/cors-proxy")
assert res.status_code == 403
def test_mcp_proxy():
global server
server.ui_mcp_proxy = True
server.start()
url = f"http://{server.server_host}:{server.server_port}/cors-proxy?url=http://example.com"
res = requests.get(url)
assert res.status_code == 200
assert "Example Domain" in res.text
def test_mcp_proxy_custom_port():
global server
server.ui_mcp_proxy = True
server.start()
# try getting the server's models API via the proxy
res = server.make_request("GET", f"/cors-proxy?url=http://{server.server_host}:{server.server_port}/models")
assert res.status_code == 200
assert "data" in res.body
def test_mcp_proxy_no_content():
# note: see issue #26598
class NoContentHandler(BaseHTTPRequestHandler):
def do_POST(self):
self.send_response(204)
self.end_headers()
def log_message(self, format, *args):
pass
target = ThreadingHTTPServer(("127.0.0.1", 0), NoContentHandler)
target_thread = threading.Thread(target=target.serve_forever, daemon=True)
target_thread.start()
try:
global server
server.ui_mcp_proxy = True
server.start()
res = server.make_request("POST", f"/cors-proxy?url=http://127.0.0.1:{target.server_port}/", data={})
assert res.status_code == 204
assert res.body in (None, b"", "")
finally:
target.shutdown()
target.server_close()
@@ -0,0 +1,171 @@
import pytest
from utils import *
server = ServerPreset.jina_reranker_tiny()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.jina_reranker_tiny()
TEST_DOCUMENTS = [
"A machine is a physical system that uses power to apply forces and control movement to perform an action. The term is commonly applied to artificial devices, such as those employing engines or motors, but also to natural biological macromolecules, such as molecular machines.",
"Learning is the process of acquiring new understanding, knowledge, behaviors, skills, values, attitudes, and preferences. The ability to learn is possessed by humans, non-human animals, and some machines; there is also evidence for some kind of learning in certain plants.",
"Machine learning is a field of study in artificial intelligence concerned with the development and study of statistical algorithms that can learn from data and generalize to unseen data, and thus perform tasks without explicit instructions.",
"Paris, capitale de la France, est une grande ville européenne et un centre mondial de l'art, de la mode, de la gastronomie et de la culture. Son paysage urbain du XIXe siècle est traversé par de larges boulevards et la Seine."
]
def test_rerank():
global server
server.start()
res = server.make_request("POST", "/rerank", data={
"query": "Machine learning is",
"documents": TEST_DOCUMENTS,
})
assert res.status_code == 200
assert len(res.body["results"]) == 4
most_relevant = res.body["results"][0]
least_relevant = res.body["results"][0]
for doc in res.body["results"]:
if doc["relevance_score"] > most_relevant["relevance_score"]:
most_relevant = doc
if doc["relevance_score"] < least_relevant["relevance_score"]:
least_relevant = doc
assert most_relevant["relevance_score"] > least_relevant["relevance_score"]
assert most_relevant["index"] == 2
assert least_relevant["index"] == 3
@pytest.mark.parametrize("kv_unified", [False, True])
def test_rerank_skips_prompt_cache(kv_unified, tmp_path):
global server
server.n_slots = 2
server.cache_ram = 100
server.kv_unified = kv_unified
server.debug = True
server.log_path = str(tmp_path / "server.log")
server.start()
for i in range(4):
res = server.make_request("POST", "/rerank", data={
"query": f"query {i}",
"documents": TEST_DOCUMENTS[i:],
})
assert res.status_code == 200
server.stop()
with open(server.log_path) as f:
log = f.read()
assert "__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" in log
assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log
def test_rerank_tei_format():
global server
server.start()
res = server.make_request("POST", "/rerank", data={
"query": "Machine learning is",
"texts": TEST_DOCUMENTS,
})
assert res.status_code == 200
assert len(res.body) == 4
most_relevant = res.body[0]
least_relevant = res.body[0]
for doc in res.body:
if doc["score"] > most_relevant["score"]:
most_relevant = doc
if doc["score"] < least_relevant["score"]:
least_relevant = doc
assert most_relevant["score"] > least_relevant["score"]
assert most_relevant["index"] == 2
assert least_relevant["index"] == 3
@pytest.mark.parametrize("documents", [
[],
None,
123,
[1, 2, 3],
])
def test_invalid_rerank_req(documents):
global server
server.start()
res = server.make_request("POST", "/rerank", data={
"query": "Machine learning is",
"documents": documents,
})
assert res.status_code == 400
assert "error" in res.body
@pytest.mark.parametrize(
"query,doc1,doc2,n_tokens",
[
("Machine learning is", "A machine", "Learning is", 19),
("Which city?", "Machine learning is ", "Paris, capitale de la", 26),
]
)
def test_rerank_usage(query, doc1, doc2, n_tokens):
global server
server.start()
res = server.make_request("POST", "/rerank", data={
"query": query,
"documents": [
doc1,
doc2,
]
})
assert res.status_code == 200
assert res.body['usage']['prompt_tokens'] == res.body['usage']['total_tokens']
assert res.body['usage']['prompt_tokens'] == n_tokens
@pytest.mark.parametrize("top_n,expected_len", [
(None, len(TEST_DOCUMENTS)), # no top_n parameter
(2, 2),
(4, 4),
(99, len(TEST_DOCUMENTS)), # higher than available docs
])
def test_rerank_top_n(top_n, expected_len):
global server
server.start()
data = {
"query": "Machine learning is",
"documents": TEST_DOCUMENTS,
}
if top_n is not None:
data["top_n"] = top_n
res = server.make_request("POST", "/rerank", data=data)
assert res.status_code == 200
assert len(res.body["results"]) == expected_len
@pytest.mark.parametrize("top_n,expected_len", [
(None, len(TEST_DOCUMENTS)), # no top_n parameter
(2, 2),
(4, 4),
(99, len(TEST_DOCUMENTS)), # higher than available docs
])
def test_rerank_tei_top_n(top_n, expected_len):
global server
server.start()
data = {
"query": "Machine learning is",
"texts": TEST_DOCUMENTS,
}
if top_n is not None:
data["top_n"] = top_n
res = server.make_request("POST", "/rerank", data=data)
assert res.status_code == 200
assert len(res.body) == expected_len
@@ -0,0 +1,865 @@
import shutil
import threading
import pytest
from utils import *
server: ServerProcess
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.router()
def test_router_props():
global server
server.models_max = 2
server.no_models_autoload = True
server.start()
res = server.make_request("GET", "/props")
assert res.status_code == 200
assert res.body["role"] == "router"
assert res.body["max_instances"] == 2
assert res.body["models_autoload"] is False
assert res.body["build_info"].startswith("b")
@pytest.mark.parametrize(
"model,success",
[
("ggml-org/tinygemma3-GGUF:Q8_0", True),
("non-existent/model", False),
]
)
def test_router_chat_completion_stream(model: str, success: bool):
global server
server.start()
content = ""
ex: ServerError | None = None
try:
res = server.make_stream_request("POST", "/chat/completions", data={
"model": model,
"max_tokens": 16,
"messages": [
{"role": "user", "content": "hello"},
],
"stream": True,
})
for data in res:
if data["choices"]:
choice = data["choices"][0]
if choice["finish_reason"] in ["stop", "length"]:
assert "content" not in choice["delta"]
else:
assert choice["finish_reason"] is None
content += choice["delta"]["content"] or ''
except ServerError as e:
ex = e
if success:
assert ex is None
assert len(content) > 0
else:
assert ex is not None
assert content == ""
def _get_model_ids(is_reload: bool, headers: dict | None = None) -> set[str]:
res = server.make_request(
"GET", "/models" + ("?reload=1" if is_reload else ""), headers=headers
)
assert res.status_code == 200
return {item["id"] for item in res.body.get("data", [])}
def _get_model_status(model_id: str, headers: dict | None = None) -> str:
res = server.make_request("GET", "/models", headers=headers)
assert res.status_code == 200
for item in res.body.get("data", []):
if item.get("id") == model_id or item.get("model") == model_id:
return item["status"]["value"]
raise AssertionError(f"Model {model_id} not found in /models response")
def _wait_for_model_status(model_id: str, desired: set[str], timeout: int = 60, headers: dict | None = None) -> str:
deadline = time.time() + timeout
last_status = None
while time.time() < deadline:
last_status = _get_model_status(model_id, headers=headers)
if last_status in desired:
return last_status
time.sleep(0.01)
raise AssertionError(
f"Timed out waiting for {model_id} to reach {desired}, last status: {last_status}"
)
def _load_model_and_wait(
model_id: str, timeout: int = 60, headers: dict | None = None
) -> None:
load_res = server.make_request(
"POST", "/models/load", data={"model": model_id}, headers=headers
)
assert load_res.status_code == 200
assert isinstance(load_res.body, dict)
assert load_res.body.get("success") is True
_wait_for_model_status(model_id, {"loaded"}, timeout=timeout, headers=headers)
def test_router_unload_model():
global server
server.start()
model_id = "ggml-org/tinygemma3-GGUF:Q8_0"
_load_model_and_wait(model_id)
unload_res = server.make_request("POST", "/models/unload", data={"model": model_id})
assert unload_res.status_code == 200
assert unload_res.body.get("success") is True
_wait_for_model_status(model_id, {"unloaded"})
def test_router_models_max_evicts_lru():
global server
server.models_max = 2
server.start()
candidate_models = [
"ggml-org/tinygemma3-GGUF:Q8_0",
"ggml-org/test-model-stories260K:F32",
"ggml-org/test-model-stories260K-infill:F32",
]
# Load only the first 2 models to fill the cache
first, second, third = candidate_models[:3]
_load_model_and_wait(first, timeout=120)
_load_model_and_wait(second, timeout=120)
# Verify both models are loaded
assert _get_model_status(first) == "loaded"
assert _get_model_status(second) == "loaded"
# Load the third model - this should trigger LRU eviction of the first model
_load_model_and_wait(third, timeout=120)
# Verify eviction: third is loaded, first was evicted
assert _get_model_status(third) == "loaded"
assert _get_model_status(first) == "unloaded"
# server_lru_sched tests (relying on LLAMA_SERVER_DEBUG_FAKE_TIMING)
MODEL_A = "ggml-org/tinygemma3-GGUF:Q8_0"
MODEL_B = "ggml-org/test-model-stories260K:F32"
MODEL_C = "ggml-org/test-model-stories260K-infill:F32"
def _tokenize(model_id: str, timeout: float | None = DEFAULT_REQUEST_TIMEOUT) -> ServerResponse:
return server.make_request(
"POST", "/tokenize", data={"model": model_id, "content": "hello world"}, timeout=timeout
)
class _Bg:
"""runs one request in a thread, keeps its result, error and finish time"""
def __init__(self, fn):
self.result = None
self.error: Exception | None = None
self.done_at: float = 0.0
self._thread = threading.Thread(target=self._run, args=(fn,), daemon=True)
def _run(self, fn):
try:
self.result = fn()
except Exception as e:
self.error = e
self.done_at = time.time()
def start(self):
self._thread.start()
return self
def join(self, timeout: int = 180):
self._thread.join(timeout)
assert not self._thread.is_alive(), "background request did not finish in time"
return self
def assert_ok(self, what: str):
assert self.error is None, f"{what} raised {self.error!r}"
assert self.result is not None and self.result.status_code == 200, \
f"{what} failed: {self.result.status_code if self.result else None} {self.result.body if self.result else None}"
def test_router_queue_does_not_evict_busy_model():
"""a request that finds no free slot waits, and the model serving a request survives it"""
global server
server.models_max = 1
server.start()
_load_model_and_wait(MODEL_A, timeout=120)
busy = _Bg(lambda: _tokenize(MODEL_A)).start()
time.sleep(0.5) # let the request reach the child and take the only slot
# no slot free and MODEL_A is busy, so this queues instead of evicting mid-request
queued = _Bg(lambda: _tokenize(MODEL_B)).start()
busy.join()
queued.join()
# had MODEL_A been evicted while serving, its own request would have died
busy.assert_ok("request against the busy model")
queued.assert_ok("queued request")
_wait_for_model_status(MODEL_B, {"loaded"}, timeout=120)
assert _get_model_status(MODEL_A) == "unloaded"
def test_router_queue_coalesces_requests_for_same_model():
"""many requests for one missing model share a slot, so only one model is given up"""
global server
server.models_max = 2
server.start()
_load_model_and_wait(MODEL_A, timeout=120)
_load_model_and_wait(MODEL_B, timeout=120)
# keep MODEL_A busy so MODEL_B is the only model that can be given up
busy = _Bg(lambda: _tokenize(MODEL_A)).start()
time.sleep(0.5)
waiters = [_Bg(lambda: _tokenize(MODEL_C)).start() for _ in range(3)]
busy.join()
for w in waiters:
w.join()
busy.assert_ok("request against the busy model")
for i, w in enumerate(waiters):
w.assert_ok(f"queued request {i}")
_wait_for_model_status(MODEL_C, {"loaded"}, timeout=120)
# one entry for 3 requests means one eviction: MODEL_B goes, MODEL_A is left alone.
# without coalescing the leftover entries still ask for a slot,
# and MODEL_A is taken too as soon as it goes idle
assert _get_model_status(MODEL_A) == "loaded"
assert _get_model_status(MODEL_B) == "unloaded"
def test_router_queue_client_disconnect_keeps_model():
"""a client that leaves while queued must not cost a running model its slot"""
global server
server.models_max = 1
server.start()
_load_model_and_wait(MODEL_A, timeout=120)
busy = _Bg(lambda: _tokenize(MODEL_A)).start()
time.sleep(0.5)
# queues behind MODEL_A, then gives up long before MODEL_A goes idle
with pytest.raises(requests.exceptions.RequestException):
_tokenize(MODEL_B, timeout=1)
busy.join()
busy.assert_ok("request against the busy model")
# nobody is waiting anymore, so MODEL_A keeps its slot
time.sleep(3)
assert _get_model_status(MODEL_A) == "loaded"
assert _get_model_status(MODEL_B) == "unloaded"
def test_router_queue_is_fifo():
"""the queue is served in arrival order"""
global server
server.models_max = 1
server.start()
_load_model_and_wait(MODEL_A, timeout=120)
busy = _Bg(lambda: _tokenize(MODEL_A)).start()
time.sleep(0.5)
first = _Bg(lambda: _tokenize(MODEL_B)).start()
time.sleep(1) # keep the arrival order unambiguous
second = _Bg(lambda: _tokenize(MODEL_C)).start()
busy.join()
first.join()
second.join()
busy.assert_ok("request against the busy model")
first.assert_ok("first queued request")
second.assert_ok("second queued request")
assert first.done_at < second.done_at, "queue was not served in arrival order"
def test_router_queue_two_waiters_share_one_eviction():
"""two requests that both find the same idle model must both be served in the end"""
global server
server.models_max = 1
server.start()
_load_model_and_wait(MODEL_A, timeout=120)
# both arrive while MODEL_A is idle, so both want its slot; only one eviction can happen
first = _Bg(lambda: _tokenize(MODEL_B)).start()
second = _Bg(lambda: _tokenize(MODEL_C)).start()
first.join(90)
second.join(90)
first.assert_ok("first queued request")
second.assert_ok("second queued request")
assert _get_model_status(MODEL_A) == "unloaded"
def test_router_no_models_autoload():
global server
server.no_models_autoload = True
server.start()
model_id = "ggml-org/tinygemma3-GGUF:Q8_0"
res = server.make_request(
"POST",
"/v1/chat/completions",
data={
"model": model_id,
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 4,
},
)
assert res.status_code == 400
assert "error" in res.body
_load_model_and_wait(model_id)
success_res = server.make_request(
"POST",
"/v1/chat/completions",
data={
"model": model_id,
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 4,
},
)
assert success_res.status_code == 200
assert "error" not in success_res.body
def test_router_api_key_required():
global server
server.api_key = "sk-router-secret"
server.start()
model_id = "ggml-org/tinygemma3-GGUF:Q8_0"
auth_headers = {"Authorization": f"Bearer {server.api_key}"}
res = server.make_request(
"POST",
"/v1/chat/completions",
data={
"model": model_id,
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 4,
},
)
assert res.status_code == 401
assert res.body.get("error", {}).get("type") == "authentication_error"
_load_model_and_wait(model_id, headers=auth_headers)
authed = server.make_request(
"POST",
"/v1/chat/completions",
headers=auth_headers,
data={
"model": model_id,
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 4,
},
)
assert authed.status_code == 200
assert "error" not in authed.body
def test_router_reload_models():
"""POST /models/reload re-reads the INI preset and updates the model list."""
global server
preset_path = os.path.join(TMP_DIR, "test_reload.ini")
# Initial preset: two models
with open(preset_path, "w") as f:
f.write(
"[model-reload-a]\n"
"hf-repo = ggml-org/test-model-stories260K\n"
"\n"
"[model-reload-b]\n"
"hf-repo = ggml-org/test-model-stories260K-infill\n"
)
server.models_preset = preset_path
server.start()
ids = _get_model_ids(is_reload=False)
assert "model-reload-a" in ids
assert "model-reload-b" in ids
# Updated preset: remove a, keep b unchanged, add c
with open(preset_path, "w") as f:
f.write(
"[model-reload-b]\n"
"hf-repo = ggml-org/test-model-stories260K-infill\n"
"\n"
"[model-reload-c]\n"
"hf-repo = ggml-org/test-model-stories260K\n"
)
try:
ids = _get_model_ids(is_reload=True)
assert "model-reload-a" not in ids, "removed model should no longer appear"
assert "model-reload-b" in ids, "unchanged model should still appear"
assert "model-reload-c" in ids, "newly added model should appear"
finally:
os.remove(preset_path)
def test_router_dedup_cache_models():
"""dedup-cache-models hides the cache entry backing a preset from GET /models"""
global server
preset_path = os.path.join(TMP_DIR, "test_dedup.ini")
main_cache_id = "ggml-org/test-model-stories260K:F32"
draft_cache_id = "ggml-org/test-model-stories260K-infill:F32"
with open(preset_path, "w") as f:
f.write(
"[model-dedup]\n"
"hf-repo = ggml-org/test-model-stories260K\n"
"spec-draft-hf = ggml-org/test-model-stories260K-infill\n"
"dedup-cache-models = 1\n"
)
server.models_preset = preset_path
server.start()
try:
ids = _get_model_ids(is_reload=False)
assert "model-dedup" in ids
assert main_cache_id not in ids, "main cache model should be hidden by dedup"
assert draft_cache_id not in ids, "draft cache model should be hidden by dedup"
# other cache models are unaffected
assert "ggml-org/tinygemma3-GGUF:Q8_0" in ids
# the hidden model is only hidden from the listing, it can still be used
res = server.make_request("POST", "/tokenize", data={"model": main_cache_id, "content": "hello"})
assert res.status_code == 200
# disabling the flag brings the cache entry back on reload
with open(preset_path, "w") as f:
f.write(
"[model-dedup]\n"
"hf-repo = ggml-org/test-model-stories260K\n"
"spec-draft-hf = ggml-org/test-model-stories260K-infill\n"
)
ids = _get_model_ids(is_reload=True)
assert main_cache_id in ids
assert draft_cache_id in ids
# the flag also works from the global section
with open(preset_path, "w") as f:
f.write(
"[*]\n"
"dedup-cache-models = 1\n"
"\n"
"[model-dedup]\n"
"hf-repo = ggml-org/test-model-stories260K\n"
"spec-draft-hf = ggml-org/test-model-stories260K-infill\n"
)
ids = _get_model_ids(is_reload=True)
assert "model-dedup" in ids
assert main_cache_id not in ids, "main cache model should be hidden by global dedup"
assert draft_cache_id not in ids, "draft cache model should be hidden by global dedup"
finally:
os.remove(preset_path)
def test_router_remote_preset():
global server
server.model_hf_repo = "ggml-org/test-preset-ci"
server.model_hf_file = None
server.offline = False
server.start()
# Should see preset models in GET /models
res = server.make_request("GET", "/models")
assert res.status_code == 200
ids = {item["id"] for item in res.body.get("data", [])}
assert "tinygemma3-preset" in ids
assert "stories260K-test" in ids
# Should be able to load a preset model
model_id = "tinygemma3-preset"
_load_model_and_wait(model_id)
MODEL_DOWNLOAD_ID = "ggml-org/test-model-router-download:F16"
MODEL_DOWNLOAD_TIMEOUT = 30
def _listen_sse(
server: ServerProcess, collected: list, stop: threading.Event, ready: threading.Event | None = None
):
"""Collect /models/sse events into `collected` until `stop` is set.
When `ready` is provided, it is set once the streaming response is open,
i.e. the server has accepted the connection and registered us as a
subscriber. Callers that trigger one-shot events (e.g. download_finished)
must wait on `ready` before acting, otherwise the event can be broadcast
before this client is subscribed and be lost.
"""
url = f"http://{server.server_host}:{server.server_port}/models/sse"
try:
with requests.get(url, stream=True, timeout=MODEL_DOWNLOAD_TIMEOUT) as resp:
if ready is not None:
ready.set()
for line_bytes in resp.iter_lines():
if stop.is_set():
break
line = line_bytes.decode("utf-8")
if line.startswith("data: "):
collected.append(json.loads(line[6:]))
except Exception:
pass
def _wait_for_sse_event(collected: list, event_type: str, model: str, timeout: int) -> bool:
deadline = time.time() + timeout
while time.time() < deadline:
if any(e.get("event") == event_type and e.get("model") == model for e in collected):
return True
time.sleep(0.01)
return False
def _own_cache(name: str) -> str:
"""An empty model cache for one test: the workers run in parallel, and a download
test must not see the files or the record another test writes in the shared cache."""
path = os.path.join(TMP_DIR, name)
shutil.rmtree(path, ignore_errors=True)
os.makedirs(path)
return path
def test_router_download_model():
"""Case 1: download a model at the model limit, verify SSE events and GET /models."""
global server
server.models_max = 1
server.start()
# Ensure the model is not present before we start
server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")
# A download worker must not consume or evict a model slot
_load_model_and_wait(MODEL_B, timeout=120)
sse_events: list = []
stop = threading.Event()
sse_ready = threading.Event()
sse_thread = threading.Thread(
target=_listen_sse, args=(server, sse_events, stop, sse_ready), daemon=True
)
sse_thread.start()
# wait for the SSE client to be subscribed before triggering the download,
# otherwise the one-shot download_finished event can be broadcast before
# this client is registered and be lost
assert sse_ready.wait(10), "SSE client failed to connect"
# Trigger the download
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID})
assert res.status_code == 200
assert res.body.get("success") is True
# Wait for download_finished SSE event
finished = _wait_for_sse_event(
sse_events, "download_finished", MODEL_DOWNLOAD_ID, MODEL_DOWNLOAD_TIMEOUT
)
stop.set()
assert finished, "Never received download_finished SSE event"
assert any(
e.get("event") == "download_progress" and e.get("model") == MODEL_DOWNLOAD_ID
for e in sse_events
), "No download_progress events received"
# Model should now appear in GET /models
ids = _get_model_ids(is_reload=False)
assert MODEL_DOWNLOAD_ID in ids, f"{MODEL_DOWNLOAD_ID} not found in /models after download"
assert _get_model_status(MODEL_B) == "loaded"
def test_router_delete_model():
"""Case 2: delete the downloaded model, verify it disappears from GET /models."""
global server
server.start()
# Ensure the model exists (download it if needed)
if MODEL_DOWNLOAD_ID not in _get_model_ids(is_reload=False):
sse_events: list = []
stop = threading.Event()
sse_ready = threading.Event()
threading.Thread(
target=_listen_sse, args=(server, sse_events, stop, sse_ready), daemon=True
).start()
# subscribe before triggering the download so the one-shot
# download_finished event is not lost (see test_router_download_model)
assert sse_ready.wait(10), "SSE client failed to connect"
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID})
assert res.status_code == 200
finished = _wait_for_sse_event(
sse_events, "download_finished", MODEL_DOWNLOAD_ID, MODEL_DOWNLOAD_TIMEOUT
)
stop.set()
assert finished, "Model did not finish downloading before delete test"
# Delete the model
del_res = server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")
assert del_res.status_code == 200
assert del_res.body.get("success") is True
# Model should no longer appear in GET /models
ids = _get_model_ids(is_reload=False)
assert MODEL_DOWNLOAD_ID not in ids, f"{MODEL_DOWNLOAD_ID} still present after deletion"
def test_router_download_sidecars():
"""Case 3: sidecars are a parameter of the download, never models of their own."""
global server
server.llama_cache = _own_cache("cache_download_sidecars")
server.start()
server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")
# an unknown sidecar kind is rejected before anything starts
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": ["lora"]})
assert res.status_code == 400
assert MODEL_DOWNLOAD_ID not in _get_model_ids(is_reload=False)
sse_events: list = []
stop = threading.Event()
sse_ready = threading.Event()
threading.Thread(
target=_listen_sse, args=(server, sse_events, stop, sse_ready), daemon=True
).start()
assert sse_ready.wait(10), "SSE client failed to connect"
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": []})
assert res.status_code == 200
finished = _wait_for_sse_event(
sse_events, "download_finished", MODEL_DOWNLOAD_ID, MODEL_DOWNLOAD_TIMEOUT
)
stop.set()
assert finished, "Never received download_finished SSE event"
# the model lists the sidecars it has on disk, and no sidecar is listed as a model
res = server.make_request("GET", "/models")
assert res.status_code == 200
entries = {item["id"]: item for item in res.body.get("data", [])}
assert entries[MODEL_DOWNLOAD_ID]["sidecars"] == []
for kind in ("mmproj", "mtp", "dflash", "eagle3", "dspark"):
assert not any(i.endswith("-" + kind) for i in entries), f"a {kind} sidecar is listed as a model"
server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")
def test_router_download_conflicts():
"""A download is refused over one on its way, and over a model that is running."""
global server
server.models_max = 2
server.llama_cache = _own_cache("cache_download_conflicts")
server.start()
server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")
# a download on its way is not started again; the entry reports the state before
# the child fetches anything, so the probe lands inside the download
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID})
assert res.status_code == 200
_wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloading"}, timeout=30)
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID})
assert res.status_code == 400
assert "is already downloading" in res.body["error"]["message"]
# a loaded model holds its files, so it goes only once it is unloaded; the
# reload that follows the download parks the entry as unloaded
_wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT)
_load_model_and_wait(MODEL_DOWNLOAD_ID, timeout=120)
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID})
assert res.status_code == 400
assert "is running, unload it first" in res.body["error"]["message"]
def test_router_delete_unknown_name():
"""A name that is neither a model nor a recorded download is not found, and the
cached tags of its repo keep their files."""
global server
server.llama_cache = _own_cache("cache_delete_unknown_name")
server.start()
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": []})
assert res.status_code == 200
_wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT)
repo = MODEL_DOWNLOAD_ID.split(":")[0]
res = server.make_request("DELETE", f"/models?model={repo}")
assert res.status_code != 200
assert "is not found" in res.body["error"]["message"]
assert MODEL_DOWNLOAD_ID in _get_model_ids(is_reload=True)
server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")
def test_router_incomplete_sidecars():
"""A cached model lists the sidecars its unfinished download has not written yet, until
they are dropped."""
global server
cache_dir = _own_cache("cache_incomplete_sidecars")
server.llama_cache = cache_dir
server.start()
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": []})
assert res.status_code == 200
_wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT)
server.stop()
# a paused mtp of the cached model, as the record of an earlier run holds it
partial_path = os.path.abspath(os.path.join(cache_dir, "test-mtp.gguf.downloadInProgress"))
with open(partial_path, "w") as f:
f.write("partial")
with open(os.path.join(cache_dir, "unfinished-downloads.json"), "w") as f:
json.dump({"downloads": [{
"model": MODEL_DOWNLOAD_ID,
"sidecars": ["mtp"],
"partials": [{"kind": "mtp", "path": partial_path}],
}]}, f)
server.start()
def entry() -> dict:
res = server.make_request("GET", "/models")
assert res.status_code == 200
return {item["id"]: item for item in res.body.get("data", [])}[MODEL_DOWNLOAD_ID]
assert entry()["sidecars"] == []
assert entry()["incomplete_sidecars"] == ["mtp"]
res = server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}&sidecar=mtp")
assert res.status_code == 200
assert res.body.get("success") is True
assert not os.path.exists(partial_path)
assert "incomplete_sidecars" not in entry()
server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")
def test_router_preset_model_sidecars():
"""A preset model reports the sidecars its files have on disk, like a cached one."""
global server
preset_path = os.path.join(TMP_DIR, "test_preset_sidecars.ini")
with open(preset_path, "w") as f:
f.write("[preset-sidecars]\nhf-repo = ggml-org/tinygemma3-GGUF:Q8_0\n")
server.models_preset = preset_path
server.start()
try:
# the mmproj of that repo lands next to the model, as -hf would pull it
res = server.make_request("POST", "/models", data={"model": MODEL_A, "sidecars": ["mmproj"]})
assert res.status_code == 200
_wait_for_model_status(MODEL_A, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT)
res = server.make_request("GET", "/models")
assert res.status_code == 200
entries = {item["id"]: item for item in res.body.get("data", [])}
assert entries["preset-sidecars"]["source"] == "preset"
assert entries["preset-sidecars"]["sidecars"] == ["mmproj"]
finally:
os.remove(preset_path)
def test_router_local_preset_listing():
"""A preset on a local file resolves from no HF repo: it lists without sidecars."""
global server
model_path = os.path.join(TMP_DIR, "test_local_preset.gguf")
preset_path = os.path.join(TMP_DIR, "test_local_preset.ini")
with open(model_path, "w") as f:
f.write("GGUF")
with open(preset_path, "w") as f:
f.write(f"[local]\nmodel = {model_path}\n")
server.models_preset = preset_path
server.start()
try:
res = server.make_request("GET", "/models")
assert res.status_code == 200
entries = {item["id"]: item for item in res.body.get("data", [])}
assert entries["local"]["source"] == "preset"
assert "sidecars" not in entries["local"]
finally:
os.remove(preset_path)
os.remove(model_path)
def test_router_download_unresolved_tag():
"""A tag that resolves no file of the repo is refused before anything starts."""
global server
server.start()
model = MODEL_DOWNLOAD_ID.split(":")[0] + ":Q9_Z"
res = server.make_request("POST", "/models", data={"model": model})
assert res.status_code == 400
assert "resolves no file to download" in res.body["error"]["message"]
assert model not in _get_model_ids(is_reload=False)
def test_router_download_drops_stale_partials():
"""Posting a recorded download again deletes the partials its new plan no longer writes."""
global server
cache_dir = _own_cache("cache_download_drops_stale_partials")
stale_path = os.path.abspath(os.path.join(cache_dir, "test-stale.gguf.downloadInProgress"))
record_path = os.path.join(cache_dir, "unfinished-downloads.json")
with open(stale_path, "w") as f:
f.write("stale")
with open(record_path, "w") as f:
json.dump({"downloads": [{
"model": MODEL_DOWNLOAD_ID,
"sidecars": ["mtp"],
"partials": [{"kind": "mtp", "path": stale_path}],
}]}, f)
server.llama_cache = cache_dir
server.start()
res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": []})
assert res.status_code == 200
assert not os.path.exists(stale_path), "the stale partial survived the new record"
_wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT)
server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")
@@ -0,0 +1,245 @@
import pytest
from openai import OpenAI
from utils import *
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
server = ServerPreset.tinyllama2()
TEST_API_KEY = "sk-this-is-the-secret-key"
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
server.api_key = TEST_API_KEY
@pytest.mark.parametrize("endpoint", ["/health"])
def test_access_public_endpoint(endpoint: str):
global server
server.start()
res = server.make_request("GET", endpoint)
assert res.status_code == 200
assert "error" not in res.body
def test_access_static_assets_without_api_key():
"""Static web UI assets should not require API key authentication (issue #21229)"""
global server
server.start()
for path in ["/", "/sw.js", "/manifest.webmanifest", "/_app/version.json"]:
res = server.make_request("GET", path)
assert res.status_code == 200, f"Expected 200 for {path}, got {res.status_code}"
@pytest.mark.parametrize("api_key", [None, "invalid-key"])
def test_incorrect_api_key(api_key: str):
global server
server.start()
res = server.make_request("POST", "/completions", data={
"prompt": "I believe the meaning of life is",
}, headers={
"Authorization": f"Bearer {api_key}" if api_key else None,
})
assert res.status_code == 401
assert "error" in res.body
assert res.body["error"]["type"] == "authentication_error"
def test_correct_api_key():
global server
server.start()
res = server.make_request("POST", "/completions", data={
"prompt": "I believe the meaning of life is",
}, headers={
"Authorization": f"Bearer {TEST_API_KEY}",
})
assert res.status_code == 200
assert "error" not in res.body
assert "content" in res.body
def test_correct_api_key_anthropic_header():
global server
server.start()
res = server.make_request("POST", "/completions", data={
"prompt": "I believe the meaning of life is",
}, headers={
"X-Api-Key": TEST_API_KEY,
})
assert res.status_code == 200
assert "error" not in res.body
assert "content" in res.body
def test_openai_library_correct_api_key():
global server
server.start()
client = OpenAI(api_key=TEST_API_KEY, base_url=f"http://{server.server_host}:{server.server_port}")
res = client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": "You are a chatbot."},
{"role": "user", "content": "What is the meaning of life?"},
],
)
assert len(res.choices) == 1
@pytest.mark.parametrize("origin,cors_header,cors_header_value", [
("localhost", "Access-Control-Allow-Origin", "localhost"),
("web.mydomain.fr", "Access-Control-Allow-Origin", "web.mydomain.fr"),
("origin", "Access-Control-Allow-Credentials", "true"),
("web.mydomain.fr", "Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS"),
("web.mydomain.fr", "Access-Control-Allow-Headers", "*"),
])
def test_cors_options(origin: str, cors_header: str, cors_header_value: str):
global server
server.start()
res = server.make_request("OPTIONS", "/completions", headers={
"Origin": origin,
"Access-Control-Request-Method": "POST",
"Access-Control-Request-Headers": "Authorization",
})
assert res.status_code == 200
assert cors_header in res.headers
assert res.headers[cors_header] == cors_header_value
@pytest.mark.parametrize("origin", [
"http://localhost",
"http://localhost:8080",
"http://127.0.0.1",
"http://127.0.0.1:3000",
"http://[::1]",
"http://[::1]:3000",
])
def test_cors_origins_localhost_reflects(origin: str):
global server
server = ServerPreset.router()
server.cors_origins = "localhost"
server.start()
res = server.make_request("OPTIONS", "/completions", headers={
"Origin": origin,
"Access-Control-Request-Method": "POST",
"Access-Control-Request-Headers": "Authorization",
})
assert res.status_code == 200
assert res.headers["Access-Control-Allow-Origin"] == origin
@pytest.mark.parametrize("origin", [
"http://web.mydomain.fr",
"http://evil.com",
"http://notlocalhost",
"http://localhost.evil.com",
])
def test_cors_origins_localhost_rejects(origin: str):
global server
server = ServerPreset.router()
server.cors_origins = "localhost"
server.start()
res = server.make_request("OPTIONS", "/completions", headers={
"Origin": origin,
"Access-Control-Request-Method": "POST",
"Access-Control-Request-Headers": "Authorization",
})
assert res.status_code == 200
assert "Access-Control-Allow-Origin" not in res.headers
def test_cors_origins_defaults_to_localhost_with_tools_enabled():
global server
server = ServerPreset.router()
server.server_tools = "all"
server.start()
res = server.make_request("OPTIONS", "/completions", headers={
"Origin": "http://localhost:8080",
"Access-Control-Request-Method": "POST",
"Access-Control-Request-Headers": "Authorization",
})
assert res.status_code == 200
assert res.headers["Access-Control-Allow-Origin"] == "http://localhost:8080"
res = server.make_request("OPTIONS", "/completions", headers={
"Origin": "http://evil.com",
"Access-Control-Request-Method": "POST",
"Access-Control-Request-Headers": "Authorization",
})
assert res.status_code == 200
assert "Access-Control-Allow-Origin" not in res.headers
def test_cors_proxy_only_forwards_explicit_proxy_headers():
class CaptureHeadersHandler(BaseHTTPRequestHandler):
def do_GET(self):
self.server.captured_headers = dict(self.headers)
self.send_response(200)
self.end_headers()
self.wfile.write(b"ok")
def log_message(self, format, *args):
pass
target = ThreadingHTTPServer(("127.0.0.1", 0), CaptureHeadersHandler)
target.captured_headers = {}
target_thread = threading.Thread(target=target.serve_forever, daemon=True)
target_thread.start()
try:
server = ServerPreset.tinyllama2()
server.api_key = TEST_API_KEY
server.ui_mcp_proxy = True
server.start()
res = server.make_request("GET", f"/cors-proxy?url=http://127.0.0.1:{target.server_port}/capture", headers={
"Authorization": f"Bearer {TEST_API_KEY}",
"Proxy-Authorization": "Basic secret",
"X-Api-Key": TEST_API_KEY,
"Cookie": "session=secret",
"x-llama-server-proxy-header-accept": "application/json",
"x-llama-server-proxy-header-authorization": "Bearer explicit",
})
assert res.status_code == 200
captured = {key.lower(): value for key, value in target.captured_headers.items()}
assert captured["accept"] == "application/json"
assert captured["authorization"] == "Bearer explicit"
assert "proxy-authorization" not in captured
assert "x-api-key" not in captured
assert "cookie" not in captured
finally:
target.shutdown()
target.server_close()
@pytest.mark.parametrize(
"media_path, image_url, success",
[
(None, "file://mtmd/test-1.jpeg", False), # disabled media path, should fail
("../../../tools", "file://mtmd/test-1.jpeg", True),
("../../../tools", "file:////mtmd//test-1.jpeg", True), # should be the same file as above
("../../../tools", "file://mtmd/notfound.jpeg", False), # non-existent file
("../../../tools", "file://../mtmd/test-1.jpeg", False), # no directory traversal
]
)
def test_local_media_file(media_path, image_url, success,):
server = ServerPreset.tinygemma3()
server.media_path = media_path
server.start()
res = server.make_request("POST", "/chat/completions", data={
"max_tokens": 1,
"messages": [
{"role": "user", "content": [
{"type": "text", "text": "test"},
{"type": "image_url", "image_url": {
"url": image_url,
}},
]},
],
})
if success:
assert res.status_code == 200
else:
assert res.status_code == 400
@@ -0,0 +1,167 @@
import pytest
import time
from utils import *
server = ServerPreset.tinyllama2()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
def is_sleeping(server: ServerProcess) -> bool:
res = server.make_request("GET", "/props")
assert res.status_code == 200
return res.body["is_sleeping"]
def wait_for_sleep(server: ServerProcess, timeout: float = 10.0):
start = time.time()
while time.time() - start < timeout:
if is_sleeping(server):
return
time.sleep(0.1)
raise TimeoutError("server did not go to sleep")
def fetch_metrics(server: ServerProcess) -> str:
res = server.make_request("GET", "/metrics")
assert res.status_code == 200
assert isinstance(res.body, str)
return res.body
def get_metric(text: str, name: str) -> float:
prefix = f"llamacpp:{name} "
values = [ln for ln in text.splitlines() if ln.startswith(prefix)]
assert len(values) == 1, f"{name} not found in metrics"
return float(values[0][len(prefix):])
def test_server_sleep():
global server
server.sleep_idle_seconds = 1
server.start()
# wait a bit so that server can go to sleep
time.sleep(2)
# make sure these endpoints are still responsive after sleep
res = server.make_request("GET", "/health")
assert res.status_code == 200
res = server.make_request("GET", "/props")
assert res.status_code == 200
assert res.body["is_sleeping"] == True
res = server.make_request("GET", "/models")
assert res.status_code == 200
assert len(res.body["data"]) == 1
assert res.body["data"][0]["id"] == server.model_alias
# make a generation request to wake up the server
res = server.make_request("POST", "/completion", data={
"n_predict": 1,
"prompt": "Hello",
})
assert res.status_code == 200
# it should no longer be sleeping
res = server.make_request("GET", "/props")
assert res.status_code == 200
assert res.body["is_sleeping"] == False
def test_server_sleep_read_only_endpoints():
global server
server.sleep_idle_seconds = 1
server.server_metrics = True
server.start()
res = server.make_request("POST", "/completion", data={
"n_predict": 4,
"prompt": "Hello",
})
assert res.status_code == 200
# the first scrape resets the throughput buckets, so that the second one reports
# the same zero rates as the snapshot taken on entering sleep
fetch_metrics(server)
metrics_awake = fetch_metrics(server)
assert get_metric(metrics_awake, "tokens_predicted_total") > 0
wait_for_sleep(server)
# during sleep, metrics are served from the snapshot taken right before sleeping
assert fetch_metrics(server) == metrics_awake
# scraping /metrics must not wake the server up
assert is_sleeping(server)
def test_server_sleep_metrics_buckets():
global server
server.sleep_idle_seconds = 1
server.server_metrics = True
server.start()
res = server.make_request("POST", "/completion", data={
"n_predict": 8,
"prompt": "Hello",
})
assert res.status_code == 200
wait_for_sleep(server)
# the first scrape reports the throughput of the last generation
assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") > 0
# nothing runs while sleeping, so the next scrapes report an empty window
assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") == 0
assert is_sleeping(server)
# waking up must not report the buckets again
res = server.make_request("POST", "/tokenize", data={"content": "Hello"})
assert res.status_code == 200
assert is_sleeping(server) == False
assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") == 0
def test_server_sleep_token_counting_wake():
global server
server.sleep_idle_seconds = 1
server.start()
wait_for_sleep(server)
assert is_sleeping(server)
res = server.make_request("POST", "/chat/completions/input_tokens", data={
"messages": [
{"role": "user", "content": "Hello world"}
]
})
assert res.status_code == 200
assert res.body["input_tokens"] > 0
assert is_sleeping(server) == False
wait_for_sleep(server)
assert is_sleeping(server)
res = server.make_request("POST", "/v1/responses/input_tokens", data={
"input": "Hello world"
})
assert res.status_code == 200
assert res.body["input_tokens"] > 0
assert is_sleeping(server) == False
wait_for_sleep(server)
assert is_sleeping(server)
res = server.make_request("POST", "/v1/messages/count_tokens", data={
"messages": [
{"role": "user", "content": "Hello world"}
]
})
assert res.status_code == 200
assert res.body["input_tokens"] > 0
assert is_sleeping(server) == False
@@ -0,0 +1,926 @@
import pytest
from utils import *
import base64
import requests
import struct
# sequence state file: magic(4) version(4) payload_size(4), then payload_size llama_token words
STATE_FILE_HEADER_SIZE = 12
server = ServerPreset.tinyllama2()
@pytest.fixture(autouse=True)
def create_server(tmp_path):
global server
server = ServerPreset.tinyllama2()
server.slot_save_path = str(tmp_path)
server.temperature = 0.0
def test_slot_save_restore():
global server
server.start()
# First prompt in slot 1 should be fully processed
res = server.make_request("POST", "/completion", data={
"prompt": "What is the capital of France?",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert match_regex("(Whiskers|Flana)+", res.body["content"])
assert res.body["timings"]["prompt_n"] == 21 # all tokens are processed
# Save state of slot 1
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": "slot1.bin",
})
assert res.status_code == 200
assert res.body["n_saved"] == 84
slot_file = os.path.join(server.slot_save_path, "slot1.bin")
assert res.body["n_written"] == os.path.getsize(slot_file)
# Since we have cache, this should only process the last tokens
res = server.make_request("POST", "/completion", data={
"prompt": "What is the capital of Germany?",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert match_regex("(Jack|said)+", res.body["content"])
assert res.body["timings"]["prompt_n"] == 6 # only different part is processed
# Loading the saved cache into slot 0
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "slot1.bin",
})
assert res.status_code == 200
assert res.body["n_restored"] == 84
assert res.body["n_read"] == os.path.getsize(slot_file)
# Since we have cache, slot 0 should only process the last tokens
res = server.make_request("POST", "/completion", data={
"prompt": "What is the capital of Germany?",
"id_slot": 0,
"cache_prompt": True,
})
assert res.status_code == 200
assert match_regex("(Jack|said)+", res.body["content"])
assert res.body["timings"]["prompt_n"] == 6 # only different part is processed
# For verification that slot 1 was not corrupted during slot 0 load, same thing should work
res = server.make_request("POST", "/completion", data={
"prompt": "What is the capital of Germany?",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert match_regex("(Jack|said)+", res.body["content"])
assert res.body["timings"]["prompt_n"] == 1
def test_slot_restore_legacy_token_list():
global server
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "What is the capital of France?",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": "slot_legacy.bin",
})
assert res.status_code == 200
assert res.body["n_saved"] == 84
# rewrite the token payload into a plain token list, as written by servers that predate the packed server_tokens format
path = os.path.join(server.slot_save_path, "slot_legacy.bin")
with open(path, "rb") as f:
data = bytearray(f.read())
# the payload written by this server starts with a packed header: LLAMA_TOKEN_NULL(4) version(4) n_tokens(4)
packed_header_size = 12
payload_size = struct.unpack_from("=I", data, STATE_FILE_HEADER_SIZE - 4)[0]
payload_end = STATE_FILE_HEADER_SIZE + payload_size * 4
n_tokens = struct.unpack_from("=I", data, STATE_FILE_HEADER_SIZE + 8)[0]
assert n_tokens == 84
tokens_start = STATE_FILE_HEADER_SIZE + packed_header_size
data = data[:STATE_FILE_HEADER_SIZE] + data[tokens_start:tokens_start + n_tokens * 4] + data[payload_end:]
struct.pack_into("=I", data, STATE_FILE_HEADER_SIZE - 4, n_tokens)
with open(path, "wb") as f:
f.write(data)
# the plain token list must restore, and the restored KV must be reusable
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "slot_legacy.bin",
})
assert res.status_code == 200
assert res.body["n_restored"] == 84
res = server.make_request("POST", "/completion", data={
"prompt": "What is the capital of Germany?",
"id_slot": 0,
"cache_prompt": True,
})
assert res.status_code == 200
assert res.body["timings"]["prompt_n"] == 6 # only the different part is processed
# the slot file carries the live draft state: a restored slot drafts like a slot that was never saved
def test_slot_restore_draft_state():
global server
server.model_draft_hf_repo = server.model_hf_repo
server.spec_type = "draft-simple"
server.start()
prompt = "What is the capital of France?"
turn = {}
for id_slot in [0, 1]:
res = server.make_request("POST", "/completion", data={
"prompt": prompt,
"n_predict": 16,
"id_slot": id_slot,
"cache_prompt": True,
})
assert res.status_code == 200
turn[id_slot] = res.body["content"]
assert turn[0] == turn[1]
res = server.make_request("POST", "/slots/0?action=save", data={
"filename": "draft_slot0.bin",
})
assert res.status_code == 200
res = server.make_request("POST", "/slots/0?action=erase")
assert res.status_code == 200
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "draft_slot0.bin",
})
assert res.status_code == 200
timings = {}
for id_slot in [0, 1]:
res = server.make_request("POST", "/completion", data={
"prompt": prompt + turn[id_slot] + " And",
"n_predict": 32,
"id_slot": id_slot,
"cache_prompt": True,
})
assert res.status_code == 200
timings[id_slot] = res.body["timings"]
assert timings[0]["prompt_n"] == timings[1]["prompt_n"]
assert timings[0]["draft_n_accepted"] > 0
assert timings[0]["draft_n"] == timings[1]["draft_n"]
assert timings[0]["draft_n_accepted"] == timings[1]["draft_n_accepted"]
def test_slot_erase():
global server
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "What is the capital of France?",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert match_regex("(Whiskers|Flana)+", res.body["content"])
assert res.body["timings"]["prompt_n"] == 21 # all tokens are processed
# erase slot 1
res = server.make_request("POST", "/slots/1?action=erase")
assert res.status_code == 200
# re-run the same prompt, it should process all tokens again
res = server.make_request("POST", "/completion", data={
"prompt": "What is the capital of France?",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert match_regex("(Whiskers|Flana)+", res.body["content"])
assert res.body["timings"]["prompt_n"] == 21 # all tokens are processed
#
# Multimodal server (mmproj loaded) slot save/restore.
#
# A pure-text slot on a multimodal server and a slot containing images must both support save/restore.
# Erase remains gated on the slot's content.
#
IMG_URL_CAT = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/91_cat.png"
IMG_URL_TRUCK = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/11_truck.png"
def _get_img_base64(url: str) -> str:
response = requests.get(url)
response.raise_for_status() # Raise an exception for bad status codes
return base64.b64encode(response.content).decode("utf-8")
@pytest.fixture
def mmproj_server():
# tinygemma3 is a small multimodal model: the mmproj is provided by the HF registry API and auto-downloaded on first run.
os.environ['LLAMA_MEDIA_MARKER'] = '<__media__>'
mm_server = ServerPreset.tinygemma3()
mm_server.slot_save_path = "./tmp"
mm_server.temperature = 0.0
return mm_server
def test_slot_save_restore_text_only_on_multimodal(mmproj_server):
server = mmproj_server
server.start()
# A pure-text prompt processed on slot 1 of a multimodal server.
res = server.make_request("POST", "/completion", data={
"prompt": "The quick brown fox jumps over the lazy dog.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
prompt_n = res.body["timings"]["prompt_n"]
assert prompt_n > 0 # all tokens are processed
# Saving a pure-text slot must succeed even though an mmproj is loaded.
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": "mm_slot1.bin",
})
assert res.status_code == 200
n_saved = res.body["n_saved"]
assert n_saved > 0 # the slot KV (prompt + generated tokens) was written
# Restore the saved state into slot 0; it must round-trip exactly.
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "mm_slot1.bin",
})
assert res.status_code == 200
assert res.body["n_restored"] == n_saved
# Prefix reuse is not checked with the default SWA cache.
res = server.make_request("POST", "/completion", data={
"prompt": "The quick brown fox jumps over the lazy dog.",
"id_slot": 0,
"cache_prompt": True,
})
assert res.status_code == 200
def test_slot_save_restore_with_image(mmproj_server):
server = mmproj_server
# Use the full SWA cache so the restored image prefix can be reused.
server.swa_full = True
server.start()
prompt_cat = {
"prompt_string": "What is this: <__media__>\n",
"multimodal_data": [_get_img_base64(IMG_URL_CAT)],
}
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 1,
"cache_prompt": True,
"prompt": prompt_cat,
})
assert res.status_code == 200
content_cat = res.body["content"]
prompt_n_full = res.body["timings"]["prompt_n"]
assert res.body["timings"]["cache_n"] == 0
assert prompt_n_full > 32 # text plus image tokens are all processed
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": "mm_slot_image.bin",
})
assert res.status_code == 200
n_saved = res.body["n_saved"]
n_written = res.body["n_written"]
assert n_saved > 0
assert n_written > 0
res = server.make_request("POST", "/slots/1?action=erase")
assert res.status_code == 200
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "mm_slot_image.bin",
})
assert res.status_code == 200
assert res.body["n_restored"] == n_saved
assert res.body["n_read"] == n_written
# a different image must not reuse the restored image tokens; only the text prefix before the image is common
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": {
"prompt_string": "What is this: <__media__>\n",
"multimodal_data": [_get_img_base64(IMG_URL_TRUCK)],
},
})
assert res.status_code == 200
cache_n = res.body["timings"]["cache_n"]
assert cache_n < 16
assert res.body["timings"]["prompt_n"] == prompt_n_full - cache_n
# restore again and resend the same image: the image tokens must be reused and greedy sampling must reproduce the original content
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "mm_slot_image.bin",
})
assert res.status_code == 200
assert res.body["n_restored"] == n_saved
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": prompt_cat,
})
assert res.status_code == 200
assert res.body["timings"]["cache_n"] == prompt_n_full - 1
assert res.body["timings"]["prompt_n"] == 1
assert res.body["content"] == content_cat
def test_slot_save_restore_with_two_images(mmproj_server):
server = mmproj_server
server.swa_full = True
server.n_ctx = 2048 # two images need more than the default 512 per slot
server.start()
prompt = {
"prompt_string": "A: <__media__> B: <__media__>\n",
"multimodal_data": [_get_img_base64(IMG_URL_CAT), _get_img_base64(IMG_URL_TRUCK)],
}
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 1,
"cache_prompt": True,
"prompt": prompt,
})
assert res.status_code == 200
prompt_n_full = res.body["timings"]["prompt_n"]
assert prompt_n_full > 64
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": "mm_slot_two_images.bin",
})
assert res.status_code == 200
n_saved = res.body["n_saved"]
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "mm_slot_two_images.bin",
})
assert res.status_code == 200
assert res.body["n_restored"] == n_saved
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": prompt,
})
assert res.status_code == 200
assert res.body["timings"]["cache_n"] == prompt_n_full - 1
assert res.body["timings"]["prompt_n"] == 1
content = res.body["content"]
res = server.make_request("POST", "/slots/1?action=restore", data={
"filename": "mm_slot_two_images.bin",
})
assert res.status_code == 200
assert res.body["n_restored"] == n_saved
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": prompt,
})
assert res.status_code == 200
assert res.body["timings"]["cache_n"] == prompt_n_full - 1
assert res.body["timings"]["prompt_n"] == 1
content = res.body["content"]
assert res.body["content"] == content
def test_slot_save_restore_with_image_across_restart(mmproj_server):
server = mmproj_server
server.swa_full = True
server.start()
prompt_cat = {
"prompt_string": "What is this: <__media__>\n",
"multimodal_data": [_get_img_base64(IMG_URL_CAT)],
}
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": prompt_cat,
})
assert res.status_code == 200
content = res.body["content"]
prompt_n_full = res.body["timings"]["prompt_n"]
res = server.make_request("POST", "/slots/0?action=save", data={
"filename": "mm_slot_restart.bin",
})
assert res.status_code == 200
n_saved = res.body["n_saved"]
# restart the server with the same model and mmproj: the saved file must restore in the new process and the image KV must be reused
server.stop()
server.start()
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "mm_slot_restart.bin",
})
assert res.status_code == 200
assert res.body["n_restored"] == n_saved
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": prompt_cat,
})
assert res.status_code == 200
assert res.body["timings"]["cache_n"] == prompt_n_full - 1
assert res.body["timings"]["prompt_n"] == 1
assert res.body["content"] == content
def test_slot_save_restore_image_payload_larger_than_context(mmproj_server):
server = mmproj_server
server.swa_full = True
server.start()
# the slot context, as the server computed it (n_ctx split across the slots)
res = server.make_request("GET", "/props")
assert res.status_code == 200
n_ctx_slot = res.body["default_generation_settings"]["n_ctx"]
# a filler token, used to grow the prompt up to the slot context
res = server.make_request("POST", "/tokenize", data={"content": " hello" * 8})
assert res.status_code == 200
assert len(res.body["tokens"]) == 8
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": {
"prompt_string": "What is this: <__media__>\n",
"multimodal_data": [_get_img_base64(IMG_URL_CAT)],
},
})
assert res.status_code == 200
prompt_cat = {
"prompt_string": "What is this: <__media__>\n" + " hello" * (n_ctx_slot - res.body["timings"]["prompt_n"] - 8),
"multimodal_data": [_get_img_base64(IMG_URL_CAT)],
}
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": prompt_cat,
})
assert res.status_code == 200
prompt_n_full = res.body["timings"]["cache_n"] + res.body["timings"]["prompt_n"]
res = server.make_request("POST", "/slots/0?action=save", data={
"filename": "mm_slot_large_payload.bin",
})
assert res.status_code == 200
path = os.path.join(server.slot_save_path, "mm_slot_large_payload.bin")
with open(path, "rb") as f:
data = bytearray(f.read())
payload_size = struct.unpack_from("=I", data, STATE_FILE_HEADER_SIZE - 4)[0]
assert payload_size > n_ctx_slot # the scenario under test: the payload does not fit in n_ctx
# drop the image from the slot, then restore it from the file
res = server.make_request("POST", "/completion", data={
"prompt": "The quick brown fox",
"id_slot": 0,
"cache_prompt": True,
})
assert res.status_code == 200
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "mm_slot_large_payload.bin",
})
assert res.status_code == 200
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": prompt_cat,
})
assert res.status_code == 200
assert res.body["timings"]["cache_n"] == prompt_n_full - 1
assert res.body["timings"]["prompt_n"] == 1
def test_slot_restore_media_file_without_mmproj(mmproj_server):
server = mmproj_server
server.start()
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": {
"prompt_string": "What is this: <__media__>\n",
"multimodal_data": [_get_img_base64(IMG_URL_CAT)],
},
})
assert res.status_code == 200
res = server.make_request("POST", "/slots/0?action=save", data={
"filename": "mm_slot_no_mmproj.bin",
})
assert res.status_code == 200
# restart the same model without the mmproj: restoring the media file must fail gracefully and leave the slot usable
server.stop()
server.no_mmproj = True
server.start()
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "mm_slot_no_mmproj.bin",
})
assert res.status_code == 400
assert "Cannot restore media tokens without an mmproj" in res.body["error"]["message"]
# A failed restore must leave the slot empty and usable.
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 1,
"cache_prompt": True,
"prompt": "The quick brown fox",
})
assert res.status_code == 200
content = res.body["content"]
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 0,
"cache_prompt": True,
"prompt": "The quick brown fox",
})
assert res.status_code == 200
assert res.body["timings"]["cache_n"] == 0
assert res.body["content"] == content
@pytest.fixture
def swa_server():
swa = ServerPreset.tinygemma3()
swa.slot_save_path = "./tmp"
swa.temperature = 0.0
swa.cache_ram = 0
# Keep the first prompt checkpoint before the divergence point.
swa.n_ubatch = 32
return swa
# the non-ASCII name checks that the appendix lands in the same file as the llama state on Windows
@pytest.mark.parametrize("filename", ["ckpt_slot1.bin", "ckpt_slot1_é.bin"])
def test_slot_restore_preserves_context_checkpoints(swa_server, filename):
server = swa_server
server.start()
base = "The quick brown fox jumps over the lazy dog. " * 20
res = server.make_request("POST", "/completion", data={
"prompt": base + "The first ending of this story is a happy one.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
n_full = res.body["timings"]["prompt_n"]
res = server.make_request("POST", "/completion", data={
"prompt": base + "But the second ending was different and sad.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
n_live = res.body["timings"]["prompt_n"]
assert n_live < n_full
res = server.make_request("POST", "/slots/1?action=erase")
assert res.status_code == 200
res = server.make_request("POST", "/completion", data={
"prompt": base + "The first ending of this story is a happy one.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": filename,
})
assert res.status_code == 200
assert res.body["n_saved"] > 0
ckpt_file = os.path.join(server.slot_save_path, filename)
assert res.body["n_written"] == os.path.getsize(ckpt_file)
res = server.make_request("POST", "/completion", data={
"prompt": "Unrelated text with no common prefix occupies the slot now.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
res = server.make_request("POST", "/slots/1?action=restore", data={
"filename": filename,
})
assert res.status_code == 200
assert res.body["n_read"] == os.path.getsize(ckpt_file)
res = server.make_request("POST", "/completion", data={
"prompt": base + "But the second ending was different and sad.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert res.body["timings"]["prompt_n"] == n_live
# checkpoint appendix: magic(4) version(4) count(4), then per checkpoint
# n_tokens(8) pos_min(4) pos_max(4) and three blobs (target, draft, speculative), each size(8) + data
def parse_ckpt_appendix(data):
off = data.find(struct.pack("<II", 0x504b4353, 1))
assert off > 0
count = struct.unpack_from("<I", data, off + 8)[0]
ckpts = []
pos = off + 12
for _ in range(count):
start = pos
pos += 16
blobs = []
for _ in range(3):
n = struct.unpack_from("<Q", data, pos)[0]
blobs.append(pos + 8)
pos += 8 + n
ckpts.append((start, pos, blobs[0]))
assert pos == len(data)
return off, ckpts
# a damaged appendix must be ignored, or its checkpoints dropped when they fail to load, without aborting the server
@pytest.mark.parametrize("damage", ["oversized_blob", "empty_target", "corrupt_state", "many_checkpoints"])
def test_slot_restore_damaged_checkpoint_appendix(swa_server, damage):
server = swa_server
server.start()
base = "The quick brown fox jumps over the lazy dog. " * 20
res = server.make_request("POST", "/completion", data={
"prompt": base + "The first ending of this story is a happy one.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
res = server.make_request("POST", "/completion", data={
"prompt": base + "But the second ending was different and sad.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
n_live = res.body["timings"]["prompt_n"]
res = server.make_request("POST", "/slots/1?action=erase")
assert res.status_code == 200
res = server.make_request("POST", "/completion", data={
"prompt": base + "The first ending of this story is a happy one.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": "ckpt_damaged.bin",
})
assert res.status_code == 200
path = os.path.join(server.slot_save_path, "ckpt_damaged.bin")
with open(path, "rb") as f:
data = bytearray(f.read())
off, ckpts = parse_ckpt_appendix(data)
if damage == "oversized_blob":
# the first target blob declares a size that cannot be allocated, it must be rejected before allocating
data = data[:ckpts[0][0] + 16] + struct.pack("<Q", 1 << 62)
elif damage == "empty_target":
# the target blobs are removed and their size set to 0, a valid save never writes an empty target state
for start, end, tgt in reversed(ckpts):
size = struct.unpack_from("<Q", data, tgt - 8)[0]
data = data[:tgt - 8] + struct.pack("<Q", 0) + data[tgt + size:]
elif damage == "corrupt_state":
# the sizes are intact, but the target states do not load
for _, _, tgt in ckpts:
struct.pack_into("<I", data, tgt, 0xdeadbeef)
else:
# more than 1024 entries: one-byte fillers that never match go first, the real checkpoints stay last
filler = struct.pack("<qiiQBQQ", 0, 0, 1 << 30, 1, 0, 0, 0)
data = data[:off + 12] + filler * (1025 - len(ckpts)) + data[off + 12:]
struct.pack_into("<I", data, off + 8, 1025)
with open(path, "wb") as f:
f.write(data)
res = server.make_request("POST", "/slots/1?action=restore", data={
"filename": "ckpt_damaged.bin",
})
assert res.status_code == 200
if damage in ("oversized_blob", "empty_target"):
assert res.body["n_read"] == off
res = server.make_request("POST", "/completion", data={
"prompt": base + "But the second ending was different and sad.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
if damage == "many_checkpoints":
assert res.body["timings"]["prompt_n"] == n_live
else:
assert res.body["timings"]["prompt_n"] > n_live
# the draft blobs of the checkpoint appendix are not covered by the main payload checks,
# so restoring into a server with another draft KV cache type must not abort
@pytest.mark.parametrize("ctkd_restore", ["f16", "q8_0"])
def test_slot_restore_checkpoints_draft_kv_type_change(swa_server, ctkd_restore):
server = swa_server
server.model_draft_hf_repo = "ggml-org/tinygemma3-GGUF:Q8_0" # same file as the target, already in the HF cache
server.spec_type = "draft-simple"
server.ctkd = "f16"
server.start()
base = "The quick brown fox jumps over the lazy dog. " * 20
res = server.make_request("POST", "/completion", data={
"prompt": base + "The first ending of this story is a happy one.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
res = server.make_request("POST", "/completion", data={
"prompt": base + "But the second ending was different and sad.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
n_live = res.body["timings"]["prompt_n"]
res = server.make_request("POST", "/slots/1?action=erase")
assert res.status_code == 200
res = server.make_request("POST", "/completion", data={
"prompt": base + "The first ending of this story is a happy one.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": "ckpt_draft_slot1.bin",
})
assert res.status_code == 200
server.stop()
server.ctkd = ctkd_restore
server.start()
res = server.make_request("POST", "/slots/1?action=restore", data={
"filename": "ckpt_draft_slot1.bin",
})
assert res.status_code == 200
res = server.make_request("POST", "/completion", data={
"prompt": base + "But the second ending was different and sad.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert res.body["timings"]["prompt_n"] == n_live
# a SWA window smaller than the prompt, so a saved state holds only the cells inside the window
@pytest.fixture
def swa_window_server(tmp_path):
swa = ServerPreset.tinygemma3()
swa.override_kv = "gemma3.attention.sliding_window=int:64"
swa.slot_save_path = str(tmp_path)
swa.temperature = 0.0
swa.n_predict = 8
return swa
SWA_PROMPT = "The quick brown fox jumps over the lazy dog. " * 30
def tokenize(server, text):
res = server.make_request("POST", "/tokenize", data={"content": text})
assert res.status_code == 200
return res.body["tokens"]
# the first turn and the exact token continuation of it, so the next turn extends the saved state
def swa_turns(server, id_slot):
res = server.make_request("POST", "/completion", data={
"prompt": tokenize(server, SWA_PROMPT),
"id_slot": id_slot,
"cache_prompt": True,
"return_tokens": True,
})
assert res.status_code == 200
return tokenize(server, SWA_PROMPT) + res.body["tokens"] + tokenize(server, " The end.")
def swa_next(server, id_slot, prompt):
res = server.make_request("POST", "/completion", data={
"prompt": prompt,
"id_slot": id_slot,
"cache_prompt": True,
})
assert res.status_code == 200
return res.body["timings"]
# a restored SWA state holds the window the next token needs, so it is reused as is
def test_slot_restore_swa_window(swa_window_server):
server = swa_window_server
server.cache_ram = 0
server.start()
prompt = {id_slot: swa_turns(server, id_slot) for id_slot in [0, 1]}
assert prompt[0] == prompt[1]
res = server.make_request("POST", "/slots/0?action=save", data={"filename": "swa_slot0.bin"})
assert res.status_code == 200
res = server.make_request("POST", "/slots/0?action=erase")
assert res.status_code == 200
res = server.make_request("POST", "/slots/0?action=restore", data={"filename": "swa_slot0.bin"})
assert res.status_code == 200
timings = {id_slot: swa_next(server, id_slot, prompt[id_slot]) for id_slot in [0, 1]}
assert timings[0]["cache_n"] == timings[1]["cache_n"]
assert timings[0]["prompt_n"] == timings[1]["prompt_n"]
# a SWA state loaded back from the RAM prompt cache is reused as is
def test_prompt_cache_swa_window(swa_window_server):
server = swa_window_server
server.n_slots = 1
def run(other_prompt_first):
server.start()
prompt = swa_turns(server, 0)
if other_prompt_first:
res = server.make_request("POST", "/completion", data={"prompt": "Once upon a time " * 60, "cache_prompt": True})
assert res.status_code == 200
timings = swa_next(server, 0, prompt)
server.stop()
return timings
direct = run(False)
cached = run(True)
assert cached["cache_n"] == direct["cache_n"]
assert cached["prompt_n"] == direct["prompt_n"]
@@ -0,0 +1,267 @@
import pytest
from utils import *
# We use a F16 MOE gguf as main model, and q4_0 as draft model
server = ServerPreset.stories15m_moe()
MODEL_DRAFT_FILE_URL = "https://huggingface.co/ggml-org/tiny-llamas/resolve/main/stories15M-q4_0.gguf"
def create_server():
global server
server = ServerPreset.stories15m_moe()
# set default values
server.model_draft = download_file(MODEL_DRAFT_FILE_URL)
server.spec_type = "draft-simple"
server.spec_draft_n_min = 4
server.spec_draft_n_max = 8
server.fa = "off"
@pytest.fixture(autouse=True)
def fixture_create_server():
return create_server()
def test_with_and_without_draft():
global server
request = {
"prompt": "I believe the meaning of life is",
"temperature": 0.2,
"top_k": 5,
"seed": 4242,
"n_predict": 16,
"return_tokens": True,
}
server.model_draft = None # disable draft model
server.spec_type = None
server.start()
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
tokens_no_draft = res.body["tokens"]
server.stop()
# create new server with draft model
create_server()
server.start()
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
tokens_draft = res.body["tokens"]
assert tokens_no_draft == tokens_draft
server.stop()
create_server()
assert server.spec_draft_n_max is not None
server.spec_synth_rates = [0.0] * server.spec_draft_n_max
server.start()
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
assert res.body["timings"]["draft_n_accepted"] == 0
assert res.body["tokens"] == tokens_no_draft
def test_different_draft_min_draft_max():
global server
test_values = [
(1, 2),
(1, 4),
(4, 8),
(4, 12),
(8, 16),
]
last_content = None
for draft_min, draft_max in test_values:
server.stop()
server.spec_draft_n_min = draft_min
server.spec_draft_n_max = draft_max
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"top_k": 1,
"n_predict": 16,
})
assert res.status_code == 200
if last_content is not None:
assert last_content == res.body["content"]
last_content = res.body["content"]
def test_synth_is_deterministic():
global server
assert server.spec_draft_n_max is not None
server.spec_synth_rates = [0.75 ** (i + 1) for i in range(server.spec_draft_n_max)]
server.start()
request = {
"prompt": "I believe the meaning of life is",
"temperature": 0.2,
"top_k": 5,
"seed": 4242,
"n_predict": 32,
}
responses = [server.make_request("POST", "/completion", data=request) for _ in range(2)]
for res in responses:
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
assert responses[0].body["timings"]["draft_n"] == responses[1].body["timings"]["draft_n"]
assert responses[0].body["timings"]["draft_n_accepted"] == responses[1].body["timings"]["draft_n_accepted"]
def test_synth_ignores_target_tokens():
global server
assert server.spec_draft_n_max is not None
server.spec_synth_rates = [1.0] * server.spec_draft_n_max
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"seed": 4242,
"n_predict": 32,
})
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
assert res.body["timings"]["draft_n_accepted"] == res.body["timings"]["draft_n"]
res = server.make_request("POST", "/completion", data={
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"seed": 4242,
"n_predict": 6,
"grammar": 'root ::= "a"{5,5}',
})
assert res.status_code == 200, res.body
res = server.make_request("POST", "/completion", data={
"prompt": "Respond with only: OK",
"temperature": 0.0,
"seed": 4242,
"n_predict": 64,
"ignore_eos": True,
})
assert res.status_code == 200, res.body
assert res.body["tokens_predicted"] == 64
assert res.body["stop_type"] == "limit"
def test_draft_n_ubatch_smaller_than_target():
global server
request = {
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"top_k": 1,
"n_predict": 16,
}
# draft with the default ubatch size of the draft context
server.start()
res_default = server.make_request("POST", "/completion", data=request)
server.stop()
assert res_default.status_code == 200
assert res_default.body["timings"]["draft_n"] > 0
# a smaller draft ubatch splits the draft batches into more ubatches, but must not change the result
server.spec_draft_n_ubatch = 8
server.start()
res_small = server.make_request("POST", "/completion", data=request)
server.stop()
assert res_small.status_code == 200
assert res_small.body["timings"]["draft_n"] > 0
assert res_small.body["tokens"] == res_default.body["tokens"]
def test_slot_ctx_not_exceeded():
global server
server.n_ctx = 256
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "Hello " * 248,
"temperature": 0.0,
"top_k": 1,
"speculative.p_min": 0.0,
})
assert res.status_code == 200
assert len(res.body["content"]) > 0
def test_with_ctx_shift():
global server
server.n_ctx = 256
server.enable_ctx_shift = True
server.start()
res = server.make_request("POST", "/completion", data={
"prompt": "Hello " * 248,
"temperature": 0.0,
"top_k": 1,
"n_predict": 256,
"speculative.p_min": 0.0,
})
assert res.status_code == 200
assert len(res.body["content"]) > 0
assert res.body["tokens_predicted"] == 256
assert res.body["truncated"] == True
@pytest.mark.parametrize("n_slots,n_requests", [
(1, 2),
(2, 2),
])
def test_multi_requests_parallel(n_slots: int, n_requests: int):
global server
server.n_slots = n_slots
server.start()
tasks = []
for _ in range(n_requests):
tasks.append((server.make_request, ("POST", "/completion", {
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"top_k": 1,
})))
results = parallel_function_calls(tasks)
for res in results:
assert res.status_code == 200
assert match_regex("(wise|kind|owl|answer)+", res.body["content"])
def test_draft_token_probs():
# tokens accepted from the draft must carry the target model's probabilities,
# same as without speculative decoding (#24271)
global server
request = {
"prompt": "I believe the meaning of life is",
"temperature": 0.0,
"top_k": 1,
"n_predict": 16,
"n_probs": 4,
}
server.model_draft = None # disable draft model
server.spec_type = None
server.start()
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
probs_no_draft = res.body["completion_probabilities"]
server.stop()
create_server()
server.start()
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
assert res.body["timings"]["draft_n"] > 0
probs_draft = res.body["completion_probabilities"]
assert len(probs_draft) == len(probs_no_draft)
for tok_draft, tok_no_draft in zip(probs_draft, probs_no_draft):
assert tok_draft["id"] == tok_no_draft["id"]
assert len(tok_draft["top_logprobs"]) == len(tok_no_draft["top_logprobs"])
assert abs(tok_draft["logprob"] - tok_no_draft["logprob"]) < 0.25
# real logprobs are not all exactly 0.0 (prob 1.0)
assert any(tok["logprob"] < -1e-4 for tok in probs_draft)
@@ -0,0 +1,153 @@
import json
import socket
import threading
import time
from urllib.parse import quote
import pytest
from utils import *
server: ServerProcess
# a model name with slashes exercises the query string routing of the stream routes: the id
# cannot travel as a path param because the decoded slash would split it before capture
MODEL = "ggml-org/tinygemma3-GGUF:Q8_0"
STREAM_ID = f"conv-stream-test::{MODEL}"
QS = "conv_id=" + quote(STREAM_ID, safe="")
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.router()
def test_stream_resume_and_stop_with_slashed_model_name():
global server
server.start()
content = ""
for data in server.make_stream_request("POST", "/chat/completions", data={
"model": MODEL,
"stream": True,
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello"}],
}, headers={"X-Conversation-Id": STREAM_ID}):
if data["choices"]:
content += data["choices"][0]["delta"].get("content") or ""
assert len(content) > 0
# the finished session replays from the beginning through the router
res = server.make_request("GET", f"/v1/stream?{QS}&from=0")
assert res.status_code == 200
assert "data: " in str(res.body)
# the explicit stop reaches the owning child and evicts the session
res = server.make_request("DELETE", f"/v1/stream?{QS}")
assert res.status_code == 204
res = server.make_request("GET", f"/v1/stream?{QS}&from=0")
assert res.status_code == 404
def test_stream_stop_during_model_load():
global server
server.start()
thread_error: list[ServerError] = []
thread_done = threading.Event()
def fire_post():
try:
for _ in server.make_stream_request("POST", "/chat/completions", data={
"model": MODEL,
"stream": True,
"max_tokens": 512,
"messages": [{"role": "user", "content": "Count from 1 to 1000."}],
}, headers={"X-Conversation-Id": STREAM_ID}):
pass
except ServerError as e:
thread_error.append(e)
finally:
thread_done.set()
t = threading.Thread(target=fire_post)
t.start()
# catch the autoload window, tiny models load fast so poll aggressively
saw_loading = False
deadline = time.time() + 5.0
while time.time() < deadline and not thread_done.is_set():
res = server.make_request("GET", "/models")
status = next(m["status"]["value"] for m in res.body["data"] if m["id"] == MODEL)
if status == "loading":
saw_loading = True
break
time.sleep(0.002)
if not saw_loading:
t.join()
pytest.skip("load window too short to be observed on this machine") # ty: ignore[too-many-positional-arguments]
# a stop during the load cancels the parked request instead of leaving an orphan
res = server.make_request("DELETE", f"/v1/stream?{QS}")
assert res.status_code == 204
assert thread_done.wait(timeout=60)
t.join()
assert len(thread_error) == 1
assert thread_error[0].code == 400
assert "cancelled" in json.dumps(thread_error[0].body)
res = server.make_request("GET", f"/v1/stream?{QS}&from=0")
assert res.status_code == 404
def test_stream_resumes_after_reload_during_model_load():
global server
server.start()
# raw socket client so the connection can be dropped mid load like a page reload
body = json.dumps({
"model": MODEL,
"stream": True,
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello"}],
})
request = (
f"POST /v1/chat/completions HTTP/1.1\r\n"
f"Host: {server.server_host}:{server.server_port}\r\n"
f"Content-Type: application/json\r\n"
f"X-Conversation-Id: {STREAM_ID}\r\n"
f"Content-Length: {len(body)}\r\n"
f"Connection: close\r\n\r\n{body}"
)
sock = socket.create_connection((server.server_host, server.server_port))
sock.sendall(request.encode())
# drop the client while the model loads, poll aggressively to catch the window
saw_loading = False
saw_503 = False
deadline = time.time() + 5.0
while time.time() < deadline:
res = server.make_request("GET", "/models")
status = next(m["status"]["value"] for m in res.body["data"] if m["id"] == MODEL)
if status == "loading":
saw_loading = True
break
if status == "loaded":
break
time.sleep(0.002)
sock.close()
if not saw_loading:
pytest.skip("load window too short to be observed on this machine") # ty: ignore[too-many-positional-arguments]
# while the model loads the resume route answers retry later, then the session appears,
# receives the whole generation despite the dead client, and replays from the beginning
deadline = time.time() + 60.0
replay = None
while time.time() < deadline:
res = server.make_request("GET", f"/v1/stream?{QS}&from=0")
if res.status_code == 503:
saw_503 = True
elif res.status_code == 200 and "data: " in str(res.body):
replay = res
break
time.sleep(0.1)
assert saw_503, "resume during the load did not answer 503"
assert replay is not None, "session never became resumable after the client disconnect"
@@ -0,0 +1,280 @@
import pytest
from utils import *
from test_vision_api import get_img_url
server = ServerPreset.tinylaya()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinylaya()
TEST_STATE = "I was charged twice for my order last week and nobody has replied."
TEST_QUESTIONS = {
"route": {
"type": "choice",
"instructions": "Which team should handle this?",
"criteria": {"billing": "payments and refunds", "shipping": None, "technical": None},
},
"urgency": {
"type": "score",
"instructions": "How urgent is this?",
"criteria": ["can wait", "this week", "today", "right now"],
},
"angry": {
"type": "noul",
"instructions": "Is the customer angry?",
},
}
def get_prompt_metrics(server: ServerProcess) -> tuple[int, int]:
"""returns the number of prompt tokens (processed, cached) since the server started"""
res = server.make_request("GET", "/metrics")
assert res.status_code == 200
values = {}
for line in res.body.splitlines():
if line.startswith("llamacpp:"):
name, value = line.split(" ")
values[name] = int(float(value))
return values["llamacpp:prompt_tokens_total"], values["llamacpp:prompt_tokens_cached_total"]
@pytest.mark.parametrize("preset", ["tinylaya", "tinyopenjev"])
def test_systemone(preset: str):
global server
server = getattr(ServerPreset, preset)()
server.start()
res = server.make_request("POST", "/v1/systemone", data={
"state": TEST_STATE,
"questions": TEST_QUESTIONS,
})
assert res.status_code == 200
assert res.body["usage"]["input_tokens"] > 0
assert res.body["usage"]["output_tokens"] == 0
answers = res.body["answers"]
assert list(answers.keys()) == ["route", "urgency", "angry"]
route = answers["route"]
assert route["type"] == "choice"
assert list(route["probabilities"].keys()) == ["billing", "shipping", "technical"]
assert abs(sum(route["probabilities"].values()) - 1.0) < 1e-4
assert route["choice"] == max(route["probabilities"], key=route["probabilities"].get)
assert 0.0 <= route["confidence"] <= 1.0
urgency = answers["urgency"]
assert urgency["type"] == "score"
assert urgency["legend"] == {"0": "can wait", "1": "this week", "2": "today", "3": "right now"}
assert list(urgency["probabilities"].keys()) == ["0", "1", "2", "3"]
assert abs(sum(urgency["probabilities"].values()) - 1.0) < 1e-4
assert abs(urgency["score"] - sum(i * p for i, p in enumerate(urgency["probabilities"].values()))) < 1e-4
assert 0.0 <= urgency["confidence"] <= 1.0
angry = answers["angry"]
assert angry["type"] == "noul"
assert 0.0 <= angry["noul"] <= 1.0
def test_systemone_json_state():
global server
server.start()
questions = {
"refund": {
"type": "noul",
"instructions": "Is a refund requested?",
"criteria": {"false": "no refund is asked", "true": "a refund is asked"},
},
}
res_obj = server.make_request("POST", "/v1/systemone", data={
"state": {"ticket": TEST_STATE, "plan": "pro"},
"questions": questions,
})
assert res_obj.status_code == 200
# an object is given to the model as JSON text
res_str = server.make_request("POST", "/v1/systemone", data={
"state": '{"ticket": "' + TEST_STATE + '", "plan": "pro"}',
"questions": questions,
})
assert res_str.status_code == 200
assert res_obj.body["usage"] == res_str.body["usage"]
assert abs(res_obj.body["answers"]["refund"]["noul"] - res_str.body["answers"]["refund"]["noul"]) < 1e-4
@pytest.mark.parametrize("data", [
{"questions": TEST_QUESTIONS},
{"state": TEST_STATE},
{"state": TEST_STATE, "questions": {}},
{"state": TEST_STATE, "questions": {"q": {"type": "unknown", "instructions": "x"}}},
{"state": TEST_STATE, "questions": {"q": {"type": "noul"}}},
{"state": TEST_STATE, "questions": {"q": {"type": "choice", "instructions": "x"}}},
{"state": TEST_STATE, "questions": {"q": {"type": "choice", "instructions": "x", "criteria": {}}}},
{"state": TEST_STATE, "questions": {"q": {"type": "score", "instructions": "x", "criteria": ["only one"]}}},
])
def test_systemone_invalid_request(data: dict):
global server
server.start()
res = server.make_request("POST", "/v1/systemone", data=data)
assert res.status_code == 400
assert "error" in res.body
def test_systemone_shared_prompt():
global server
server = ServerPreset.tinyopenjev()
server.server_metrics = True
server.start()
res = server.make_request("POST", "/v1/systemone", data={
"state": TEST_STATE,
"questions": TEST_QUESTIONS,
})
assert res.status_code == 200
# the first question evaluates the shared prefix, the 2 others start from it
n_processed, n_cached = get_prompt_metrics(server)
assert n_cached > 0
assert n_cached % 2 == 0
assert n_processed + n_cached == res.body["usage"]["input_tokens"]
# with one slot the prompt cannot be shared, the answers must be the same
server.stop()
server = ServerPreset.tinyopenjev()
server.n_slots = 1
server.start()
res_single = server.make_request("POST", "/v1/systemone", data={
"state": TEST_STATE,
"questions": TEST_QUESTIONS,
})
assert res_single.status_code == 200
assert res_single.body["usage"] == res.body["usage"]
for qid in ["route", "urgency"]:
probs_shared = res.body["answers"][qid]["probabilities"]
probs_single = res_single.body["answers"][qid]["probabilities"]
for key in probs_shared:
assert abs(probs_shared[key] - probs_single[key]) < 0.01
assert abs(res.body["answers"]["angry"]["noul"] - res_single.body["answers"]["angry"]["noul"]) < 0.01
@pytest.mark.parametrize("cache_ram", [None, 100])
def test_systemone_prompt_cache(cache_ram, tmp_path):
global server
server.n_slots = 1
server.cache_ram = cache_ram
server.debug = True
server.log_path = str(tmp_path / "server.log")
server.start()
# the third request comes back to the state of the first one
for state in ["Ticket A: " + TEST_STATE, "Ticket B: " + TEST_STATE * 2, "Ticket A: " + TEST_STATE]:
res = server.make_request("POST", "/v1/systemone", data={
"state": state,
"questions": TEST_QUESTIONS,
})
assert res.status_code == 200
server.stop()
with open(server.log_path) as f:
log = f.read()
# the prompt cache is opt-in for decision models, and restores their states once enabled
assert ("__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" in log) == (cache_ram is not None)
assert ("found better prompt" in log) == (cache_ram is not None)
def test_systemone_prompt_cache_child_slots(tmp_path):
global server
server = ServerPreset.tinyopenjev()
server.n_slots = len(TEST_QUESTIONS)
server.cache_ram = 100
server.debug = True
server.log_path = str(tmp_path / "server.log")
server.start()
states = [
"My parcel shows delivered but it never arrived at my door.",
"The app crashes every time I open the settings page.",
"Please close my account and delete all my personal data.",
]
angry = {"angry": TEST_QUESTIONS["angry"]}
def decide(state, questions):
res = server.make_request("POST", "/v1/systemone", data={"state": state, "questions": questions})
assert res.status_code == 200
# one state per slot, then a request whose child tasks take every slot
for state in states:
decide(state, angry)
decide(TEST_STATE, TEST_QUESTIONS)
with open(server.log_path) as f:
n_restored = f.read().count("found better prompt with f_keep = 1.000")
# every overwritten state comes back whole from the prompt cache
for state in states:
decide(state, angry)
server.stop()
with open(server.log_path) as f:
assert f.read().count("found better prompt with f_keep = 1.000") == n_restored + len(states)
def test_systemone_images():
global server
server = ServerPreset.tinyopenjev()
server.start()
image = get_img_url("IMG_BASE64_URI_0")
res_text = server.make_request("POST", "/v1/systemone", data={
"state": TEST_STATE,
"questions": TEST_QUESTIONS,
})
assert res_text.status_code == 200
res = server.make_request("POST", "/v1/systemone", data={
"state": TEST_STATE,
"questions": TEST_QUESTIONS,
"images": [image],
})
assert res.status_code == 200
assert list(res.body["answers"].keys()) == ["route", "urgency", "angry"]
assert res.body["usage"]["input_tokens"] > res_text.body["usage"]["input_tokens"]
# same image, given as a part of a chat message
res_part = server.make_request("POST", "/v1/systemone", data={
"state": [{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": image}},
{"type": "text", "text": TEST_STATE},
]}],
"questions": TEST_QUESTIONS,
})
assert res_part.status_code == 200
assert res_part.body["usage"]["input_tokens"] > res_text.body["usage"]["input_tokens"]
res = server.make_request("POST", "/v1/systemone", data={
"state": TEST_STATE,
"questions": TEST_QUESTIONS,
"images": [image] * 9,
})
assert res.status_code == 400
res = server.make_request("POST", "/v1/systemone", data={
"state": TEST_STATE,
"questions": TEST_QUESTIONS,
"images": ["https://example.com/image.png"],
})
assert res.status_code == 400
def test_systemone_images_not_supported():
global server
server.start()
res = server.make_request("POST", "/v1/systemone", data={
"state": TEST_STATE,
"questions": TEST_QUESTIONS,
"images": [get_img_url("IMG_BASE64_URI_0")],
})
assert res.status_code == 501
@@ -0,0 +1,106 @@
#!/usr/bin/env python
import pytest
# ensure grandparent path is in sys.path
from pathlib import Path
import sys
from unit.test_tool_call import TEST_TOOL
path = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(path))
import datetime
from utils import *
from typing import Literal
server: ServerProcess
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
server.model_alias = "tinyllama-2"
server.n_slots = 1
@pytest.mark.parametrize("tools", [None, [], [TEST_TOOL]])
@pytest.mark.parametrize("template_name,reasoning,expected_end", [
("deepseek-ai-DeepSeek-R1-Distill-Qwen-32B", "on", "<think>\n"),
("deepseek-ai-DeepSeek-R1-Distill-Qwen-32B","auto", "<think>\n"),
("deepseek-ai-DeepSeek-R1-Distill-Qwen-32B", "off", "<think>\n</think>"),
("Qwen-Qwen3-0.6B","auto", "<|im_start|>assistant\n"),
("Qwen-Qwen3-0.6B", "off", "<|im_start|>assistant\n<think>\n\n</think>\n\n"),
("Qwen-QwQ-32B","auto", "<|im_start|>assistant\n<think>\n"),
("Qwen-QwQ-32B", "off", "<|im_start|>assistant\n<think>\n</think>"),
("CohereForAI-c4ai-command-r7b-12-2024-tool_use","auto", "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>"),
("CohereForAI-c4ai-command-r7b-12-2024-tool_use", "off", "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|><|END_THINKING|>"),
])
def test_reasoning(template_name: str, reasoning: Literal['on', 'off', 'auto'] | None, expected_end: str, tools: list[dict]):
global server
server.jinja = True
server.reasoning = reasoning
server.chat_template_file = f'../../../models/templates/{template_name}.jinja'
server.start()
res = server.make_request("POST", "/apply-template", data={
"messages": [
{"role": "user", "content": "What is today?"},
],
"tools": tools,
})
assert res.status_code == 200
prompt = res.body["prompt"]
assert prompt.endswith(expected_end), f"Expected prompt to end with '{expected_end}', got '{prompt}'"
@pytest.mark.parametrize("tools", [None, [], [TEST_TOOL]])
@pytest.mark.parametrize("template_name,format", [
("meta-llama-Llama-3.3-70B-Instruct", "%d %b %Y"),
("fireworks-ai-llama-3-firefunction-v2", "%b %d %Y"),
])
def test_date_inside_prompt(template_name: str, format: str, tools: list[dict]):
global server
server.jinja = True
server.chat_template_file = f'../../../models/templates/{template_name}.jinja'
server.start()
res = server.make_request("POST", "/apply-template", data={
"messages": [
{"role": "user", "content": "What is today?"},
],
"tools": tools,
})
assert res.status_code == 200
prompt = res.body["prompt"]
today_str = datetime.date.today().strftime(format)
assert today_str in prompt, f"Expected today's date ({today_str}) in content ({prompt})"
@pytest.mark.parametrize("add_generation_prompt", [False, True])
@pytest.mark.parametrize("template_name,expected_generation_prompt", [
("meta-llama-Llama-3.3-70B-Instruct", "<|start_header_id|>assistant<|end_header_id|>"),
])
def test_add_generation_prompt(template_name: str, expected_generation_prompt: str, add_generation_prompt: bool):
global server
server.jinja = True
server.chat_template_file = f'../../../models/templates/{template_name}.jinja'
server.start()
res = server.make_request("POST", "/apply-template", data={
"messages": [
{"role": "user", "content": "What is today?"},
],
"add_generation_prompt": add_generation_prompt,
})
assert res.status_code == 200
prompt = res.body["prompt"]
if add_generation_prompt:
assert expected_generation_prompt in prompt, f"Expected generation prompt ({expected_generation_prompt}) in content ({prompt})"
else:
assert expected_generation_prompt not in prompt, f"Did not expect generation prompt ({expected_generation_prompt}) in content ({prompt})"
@@ -0,0 +1,59 @@
import pytest
from utils import *
server = ServerPreset.tinyllama2()
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
def test_tokenize_detokenize():
global server
server.start()
# tokenize
content = "What is the capital of France ?"
res_tok = server.make_request("POST", "/tokenize", data={
"content": content
})
assert res_tok.status_code == 200
assert len(res_tok.body["tokens"]) > 5
# detokenize
res_detok = server.make_request("POST", "/detokenize", data={
"tokens": res_tok.body["tokens"],
})
assert res_detok.status_code == 200
assert res_detok.body["content"].strip() == content
def test_tokenize_with_bos():
global server
server.start()
# tokenize
content = "What is the capital of France ?"
bosId = 1
res_tok = server.make_request("POST", "/tokenize", data={
"content": content,
"add_special": True,
})
assert res_tok.status_code == 200
assert res_tok.body["tokens"][0] == bosId
def test_tokenize_with_pieces():
global server
server.start()
# tokenize
content = "This is a test string with unicode 媽 and emoji 🤗"
res_tok = server.make_request("POST", "/tokenize", data={
"content": content,
"with_pieces": True,
})
assert res_tok.status_code == 200
for token in res_tok.body["tokens"]:
assert "id" in token
assert token["id"] > 0
assert "piece" in token
assert len(token["piece"]) > 0
+644
View File
@@ -0,0 +1,644 @@
#!/usr/bin/env python
import pytest
# ensure grandparent path is in sys.path
from pathlib import Path
import sys
path = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(path))
from utils import *
from enum import Enum
from typing import TypedDict
server: ServerProcess
TIMEOUT_START_SLOW = 15 * 60 # this is needed for real model tests
TIMEOUT_HTTP_REQUEST = 60
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
server.model_alias = "tinyllama-2-tool-call"
server.n_slots = 1
server.n_ctx = 8192
server.n_batch = 2048
class CompletionMode(Enum):
NORMAL = "normal"
STREAMED = "streamed"
class ToolParameters(TypedDict):
type: str
properties: dict[str, dict]
required: list[str]
class ToolFunction(TypedDict):
name: str
description: str
parameters: ToolParameters
class ToolDefinition(TypedDict):
type: str
function: ToolFunction
TEST_TOOL = ToolDefinition(
type = "function",
function = ToolFunction(
name = "test",
description = "",
parameters = ToolParameters(
type = "object",
properties = {
"success": {
"type": "boolean",
"const": True,
},
},
required = ["success"],
),
),
)
PYTHON_TOOL = ToolDefinition(
type = "function",
function = ToolFunction(
name = "python",
description = "Runs code in an ipython interpreter and returns the result of the execution after 60 seconds.",
parameters = ToolParameters(
type = "object",
properties = {
"code": {
"type": "string",
"description": "The code to run in the ipython interpreter.",
},
},
required = ["code"],
),
),
)
WEATHER_TOOL = ToolDefinition(
type = "function",
function = ToolFunction(
name = "get_current_weather",
description = "Get the current weather in a given location",
parameters = ToolParameters(
type = "object",
properties = {
"location": {
"type": "string",
"description": "The city and country/state, e.g. 'San Francisco, CA', or 'Paris, France'",
},
},
required = ["location"],
),
),
)
def do_test_completion_with_required_tool_tiny(server: ServerProcess, tool: dict, argument_key: str | None, n_predict, **kwargs):
body = server.make_any_request("POST", "/v1/chat/completions", data={
"max_tokens": n_predict,
"messages": [
{"role": "system", "content": "You are a coding assistant."},
{"role": "user", "content": "Write an example"},
],
"tool_choice": "required",
"tools": [tool],
"parallel_tool_calls": False,
**kwargs,
})
# assert res.status_code == 200, f"Expected status code 200, got {res.status_code}"
choice = body["choices"][0]
tool_calls = choice["message"].get("tool_calls")
assert tool_calls and len(tool_calls) == 1, f'Expected 1 tool call in {choice["message"]}'
tool_call = tool_calls[0]
assert choice["message"].get("content") in (None, ""), f'Expected no content in {choice["message"]}'
# assert len(tool_call.get("id", "")) > 0, f'Expected non empty tool call id in {tool_call}'
expected_function_name = "python" if tool["type"] == "code_interpreter" else tool["function"]["name"]
assert expected_function_name == tool_call["function"]["name"], f'Expected tool name to be {tool_call["function"]["name"]} in {choice["message"]}'
actual_arguments = tool_call["function"]["arguments"]
assert isinstance(actual_arguments, dict) or isinstance(actual_arguments, str), f'Expected arguments to be a dict or str, got: {actual_arguments}'
if argument_key is not None:
if (isinstance(actual_arguments, str)):
actual_arguments = json.loads(actual_arguments)
assert argument_key in actual_arguments, f"tool arguments: {actual_arguments}, expected: {argument_key}"
# PR #22654: commented out since we're now allowing content before tool calls in tool_call: required, so we can't force this
# in the tiny model just by using the grammar
#
# @pytest.mark.parametrize("stream", [CompletionMode.NORMAL, CompletionMode.STREAMED])
# @pytest.mark.parametrize("template_name,tool,argument_key", [
# ("Qwen3-Coder", TEST_TOOL, "success"),
# ("Qwen3-Coder", TEST_TOOL, "success"),
# ("meta-llama-Llama-3.3-70B-Instruct", TEST_TOOL, "success"),
# ("meta-llama-Llama-3.3-70B-Instruct", TEST_TOOL, "success"),
# ("meta-llama-Llama-3.3-70B-Instruct", PYTHON_TOOL, "code"),
# ("meta-llama-Llama-3.3-70B-Instruct", PYTHON_TOOL, "code"),
# ])
# def test_completion_with_required_tool_tiny_fast(template_name: str, tool: dict, argument_key: str | None, stream: CompletionMode):
# global server
# n_predict = 1024
# # server = ServerPreset.stories15m_moe()
# server.jinja = True
# server.n_predict = n_predict
# server.chat_template_file = f'../../../models/templates/{template_name}.jinja'
# server.start()
# do_test_completion_with_required_tool_tiny(server, tool, argument_key, n_predict, stream=stream == CompletionMode.STREAMED, temperature=0.0, top_k=1, top_p=1.0)
# @pytest.mark.slow
# @pytest.mark.parametrize("stream", [CompletionMode.NORMAL, CompletionMode.STREAMED])
# @pytest.mark.parametrize("template_name,tool,argument_key", [
# ("meta-llama-Llama-3.1-8B-Instruct", TEST_TOOL, "success"),
# ("meta-llama-Llama-3.1-8B-Instruct", PYTHON_TOOL, "code"),
# ("meetkai-functionary-medium-v3.1", TEST_TOOL, "success"),
# ("meetkai-functionary-medium-v3.1", PYTHON_TOOL, "code"),
# ("meetkai-functionary-medium-v3.2", TEST_TOOL, "success"),
# # Functionary v3.2 format supports raw python content, which w/ a dummy stories model will never end on its own.
# # ("meetkai-functionary-medium-v3.2", PYTHON_TOOL, "code"),
# ("NousResearch-Hermes-2-Pro-Llama-3-8B-tool_use", TEST_TOOL, "success"),
# ("NousResearch-Hermes-2-Pro-Llama-3-8B-tool_use", PYTHON_TOOL, "code"),
# ("meta-llama-Llama-3.2-3B-Instruct", TEST_TOOL, "success"),
# ("meta-llama-Llama-3.2-3B-Instruct", PYTHON_TOOL, "code"),
# ("mistralai-Mistral-Nemo-Instruct-2407", TEST_TOOL, "success"),
# ("mistralai-Mistral-Nemo-Instruct-2407", PYTHON_TOOL, "code"),
# ("NousResearch-Hermes-3-Llama-3.1-8B-tool_use", TEST_TOOL, "success"),
# ("NousResearch-Hermes-3-Llama-3.1-8B-tool_use", PYTHON_TOOL, "code"),
# ("deepseek-ai-DeepSeek-R1-Distill-Llama-8B", TEST_TOOL, "success"),
# ("deepseek-ai-DeepSeek-R1-Distill-Llama-8B", PYTHON_TOOL, "code"),
# ("fireworks-ai-llama-3-firefunction-v2", TEST_TOOL, "success"),
# # ("fireworks-ai-llama-3-firefunction-v2", PYTHON_TOOL, "codeFalse), True),
# # ("fireworks-ai-llama-3-firefunction-v2", PYTHON_TOOL, "code"),
# ])
# def test_completion_with_required_tool_tiny_slow(template_name: str, tool: dict, argument_key: str | None, stream: CompletionMode):
# global server
# n_predict = 512
# # server = ServerPreset.stories15m_moe()
# server.jinja = True
# server.n_predict = n_predict
# server.chat_template_file = f'../../../models/templates/{template_name}.jinja'
# server.start(timeout_seconds=TIMEOUT_START_SLOW)
# do_test_completion_with_required_tool_tiny(server, tool, argument_key, n_predict, stream=stream == CompletionMode.STREAMED)
@pytest.mark.slow
@pytest.mark.parametrize("stream", [CompletionMode.NORMAL, CompletionMode.STREAMED])
@pytest.mark.parametrize("tool,argument_key,hf_repo,template_override", [
(TEST_TOOL, "success", "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/gemma-2-2b-it-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/gemma-2-2b-it-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/gemma-2-2b-it-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/Qwen2.5-Coder-3B-Instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Qwen2.5-Coder-3B-Instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Qwen2.5-Coder-3B-Instruct-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/Hermes-2-Pro-Llama-3-8B-GGUF:Q4_K_M", ("NousResearch/Hermes-2-Pro-Llama-3-8B", "tool_use")),
(PYTHON_TOOL, "code", "bartowski/Hermes-2-Pro-Llama-3-8B-GGUF:Q4_K_M", ("NousResearch/Hermes-2-Pro-Llama-3-8B", "tool_use")),
(PYTHON_TOOL, "code", "bartowski/Hermes-2-Pro-Llama-3-8B-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/Hermes-3-Llama-3.1-8B-GGUF:Q4_K_M", ("NousResearch/Hermes-3-Llama-3.1-8B", "tool_use")),
(PYTHON_TOOL, "code", "bartowski/Hermes-3-Llama-3.1-8B-GGUF:Q4_K_M", ("NousResearch/Hermes-3-Llama-3.1-8B", "tool_use")),
(PYTHON_TOOL, "code", "bartowski/Hermes-3-Llama-3.1-8B-GGUF:Q4_K_M", "chatml"),
# (TEST_TOOL, "success", "bartowski/Mistral-Nemo-Instruct-2407-GGUF:Q4_K_M", None),
# (PYTHON_TOOL, "code", "bartowski/Mistral-Nemo-Instruct-2407-GGUF:Q4_K_M", None),
# (PYTHON_TOOL, "code", "bartowski/Mistral-Nemo-Instruct-2407-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/functionary-small-v3.2-GGUF:Q4_K_M", ("meetkai/functionary-medium-v3.2", None)),
(PYTHON_TOOL, "code", "bartowski/functionary-small-v3.2-GGUF:Q4_K_M", ("meetkai/functionary-medium-v3.2", None)),
(PYTHON_TOOL, "code", "bartowski/functionary-small-v3.2-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M", ("meta-llama/Llama-3.2-3B-Instruct", None)),
(PYTHON_TOOL, "code", "bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M", ("meta-llama/Llama-3.2-3B-Instruct", None)),
(PYTHON_TOOL, "code", "bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M", ("meta-llama/Llama-3.2-3B-Instruct", None)),
(PYTHON_TOOL, "code", "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M", ("meta-llama/Llama-3.2-3B-Instruct", None)),
(PYTHON_TOOL, "code", "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M", "chatml"),
(TEST_TOOL, "success", "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF:Q4_K_M", None),
(PYTHON_TOOL, "code", "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF:Q4_K_M", None),
])
def test_completion_with_required_tool_real_model(tool: dict, argument_key: str | None, hf_repo: str, template_override: str | Tuple[str, str | None] | None, stream: CompletionMode):
global server
n_predict = 512
server.jinja = True
server.n_ctx = 8192
server.n_predict = n_predict
server.model_hf_repo = hf_repo
server.model_hf_file = None
if isinstance(template_override, tuple):
(template_hf_repo, template_variant) = template_override
server.chat_template_file = f"../../../models/templates/{template_hf_repo.replace('/', '-') + ('-' + template_variant if template_variant else '')}.jinja"
assert os.path.exists(server.chat_template_file), f"Template file {server.chat_template_file} does not exist. Run `python scripts/get_chat_template.py {template_hf_repo} {template_variant} > {server.chat_template_file}` to download the template."
elif isinstance(template_override, str):
server.chat_template = template_override
server.start(timeout_seconds=TIMEOUT_START_SLOW)
body = server.make_any_request("POST", "/v1/chat/completions", data={
"max_tokens": n_predict,
"messages": [
{"role": "system", "content": "You are a coding assistant."},
{"role": "user", "content": "Write an example"},
],
"tool_choice": "required",
"tools": [tool],
"parallel_tool_calls": False,
"stream": stream == CompletionMode.STREAMED,
"temperature": 0.0,
"top_k": 1,
"top_p": 1.0,
}, timeout=TIMEOUT_HTTP_REQUEST)
choice = body["choices"][0]
tool_calls = choice["message"].get("tool_calls")
assert tool_calls and len(tool_calls) == 1, f'Expected 1 tool call in {choice["message"]}'
tool_call = tool_calls[0]
# assert choice["message"].get("content") in (None, ""), f'Expected no content in {choice["message"]}'
expected_function_name = "python" if tool["type"] == "code_interpreter" else tool["function"]["name"]
assert expected_function_name == tool_call["function"]["name"]
actual_arguments = tool_call["function"]["arguments"]
assert isinstance(actual_arguments, str)
if argument_key is not None:
actual_arguments = json.loads(actual_arguments)
assert argument_key in actual_arguments, f"tool arguments: {json.dumps(actual_arguments)}, expected: {argument_key}"
def do_test_completion_without_tool_call(server: ServerProcess, n_predict: int, tools: list[dict], tool_choice: str | None, **kwargs):
body = server.make_any_request("POST", "/v1/chat/completions", data={
"max_tokens": n_predict,
"messages": [
{"role": "system", "content": "You are a coding assistant."},
{"role": "user", "content": "say hello world with python"},
],
"tools": tools if tools else None,
"tool_choice": tool_choice,
**kwargs,
}, timeout=TIMEOUT_HTTP_REQUEST)
choice = body["choices"][0]
assert choice["message"].get("tool_calls") is None, f'Expected no tool call in {choice["message"]}'
@pytest.mark.parametrize("stream", [CompletionMode.NORMAL, CompletionMode.STREAMED])
@pytest.mark.parametrize("template_name,n_predict,tools,tool_choice", [
("meta-llama-Llama-3.3-70B-Instruct", 128, [], None),
("meta-llama-Llama-3.3-70B-Instruct", 128, [TEST_TOOL], None),
("meta-llama-Llama-3.3-70B-Instruct", 128, [PYTHON_TOOL], 'none'),
])
def test_completion_without_tool_call_fast(template_name: str, n_predict: int, tools: list[dict], tool_choice: str | None, stream: CompletionMode):
global server
server.n_predict = n_predict
server.jinja = True
server.chat_template_file = f'../../../models/templates/{template_name}.jinja'
server.start()
do_test_completion_without_tool_call(server, n_predict, tools, tool_choice, stream=stream == CompletionMode.STREAMED)
@pytest.mark.slow
@pytest.mark.parametrize("stream", [CompletionMode.NORMAL, CompletionMode.STREAMED])
@pytest.mark.parametrize("template_name,n_predict,tools,tool_choice", [
("meetkai-functionary-medium-v3.2", 256, [], None),
("meetkai-functionary-medium-v3.2", 256, [TEST_TOOL], None),
("meetkai-functionary-medium-v3.2", 256, [PYTHON_TOOL], 'none'),
("meetkai-functionary-medium-v3.1", 256, [], None),
("meetkai-functionary-medium-v3.1", 256, [TEST_TOOL], None),
("meetkai-functionary-medium-v3.1", 256, [PYTHON_TOOL], 'none'),
("meta-llama-Llama-3.2-3B-Instruct", 256, [], None),
("meta-llama-Llama-3.2-3B-Instruct", 256, [TEST_TOOL], None),
("meta-llama-Llama-3.2-3B-Instruct", 256, [PYTHON_TOOL], 'none'),
])
def test_completion_without_tool_call_slow(template_name: str, n_predict: int, tools: list[dict], tool_choice: str | None, stream: CompletionMode):
global server
server.n_predict = n_predict
server.jinja = True
server.chat_template_file = f'../../../models/templates/{template_name}.jinja'
server.start(timeout_seconds=TIMEOUT_START_SLOW)
do_test_completion_without_tool_call(server, n_predict, tools, tool_choice, stream=stream == CompletionMode.STREAMED)
@pytest.mark.slow
@pytest.mark.parametrize("stream", [CompletionMode.NORMAL, CompletionMode.STREAMED])
@pytest.mark.parametrize("hf_repo,template_override", [
("bartowski/Meta-Llama-3.1-8B-Instruct-GGUF:Q4_K_M", None),
("bartowski/Meta-Llama-3.1-8B-Instruct-GGUF:Q4_K_M", "chatml"),
("bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", None),
("bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", "chatml"),
("bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", None),
("bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", "chatml"),
("bartowski/Qwen2.5-Coder-3B-Instruct-GGUF:Q4_K_M", None),
("bartowski/Qwen2.5-Coder-3B-Instruct-GGUF:Q4_K_M", "chatml"),
("bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", None),
("bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", "chatml"),
("bartowski/Hermes-2-Pro-Llama-3-8B-GGUF:Q4_K_M", ("NousResearch/Hermes-2-Pro-Llama-3-8B", "tool_use")),
("bartowski/Hermes-2-Pro-Llama-3-8B-GGUF:Q4_K_M", "chatml"),
("bartowski/Hermes-3-Llama-3.1-8B-GGUF:Q4_K_M", ("NousResearch/Hermes-3-Llama-3.1-8B", "tool_use")),
("bartowski/Hermes-3-Llama-3.1-8B-GGUF:Q4_K_M", "chatml"),
# ("bartowski/Mistral-Nemo-Instruct-2407-GGUF:Q4_K_M", None),
# ("bartowski/Mistral-Nemo-Instruct-2407-GGUF:Q4_K_M", "chatml"),
# ("bartowski/functionary-small-v3.2-GGUF:Q8_0", ("meetkai/functionary-medium-v3.2", None)),
# ("bartowski/functionary-small-v3.2-GGUF:Q8_0", "chatml"),
("bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M", ("meta-llama/Llama-3.2-3B-Instruct", None)),
("bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M", "chatml"),
("bartowski/c4ai-command-r7b-12-2024-GGUF:Q6_K_L", ("CohereForAI/c4ai-command-r7b-12-2024", "tool_use")),
("bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF:Q4_K_M", None),
# Note: gemma-2-2b-it knows itself as "model", not "assistant", so we don't test the ill-suited chatml on it.
("bartowski/gemma-2-2b-it-GGUF:Q4_K_M", None),
# ("bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M", ("meta-llama/Llama-3.2-3B-Instruct", None)),
])
def test_weather(hf_repo: str, template_override: str | Tuple[str, str | None] | None, stream: CompletionMode):
global server
n_predict = 512
server.jinja = True
server.n_ctx = 8192
server.n_predict = n_predict
server.model_hf_repo = hf_repo
server.model_hf_file = None
if isinstance(template_override, tuple):
(template_hf_repo, template_variant) = template_override
server.chat_template_file = f"../../../models/templates/{template_hf_repo.replace('/', '-') + ('-' + template_variant if template_variant else '')}.jinja"
assert os.path.exists(server.chat_template_file), f"Template file {server.chat_template_file} does not exist. Run `python scripts/get_chat_template.py {template_hf_repo} {template_variant} > {server.chat_template_file}` to download the template."
elif isinstance(template_override, str):
server.chat_template = template_override
server.start()
do_test_weather(server, stream=stream == CompletionMode.STREAMED, max_tokens=n_predict)
def do_test_weather(server: ServerProcess, **kwargs):
body = server.make_any_request("POST", "/v1/chat/completions", data={
"messages": [
{"role": "system", "content": "You are a chatbot that uses tools/functions. Dont overthink things."},
{"role": "user", "content": "What is the weather in Istanbul?"},
],
"tools": [WEATHER_TOOL],
**kwargs,
}, timeout=TIMEOUT_HTTP_REQUEST)
choice = body["choices"][0]
tool_calls = choice["message"].get("tool_calls")
assert tool_calls and len(tool_calls) == 1, f'Expected 1 tool call in {choice["message"]}'
tool_call = tool_calls[0]
# assert choice["message"].get("content") in (None, ""), f'Expected no content in {choice["message"]}'
assert tool_call["function"]["name"] == WEATHER_TOOL["function"]["name"], f'Expected weather tool call, got {tool_call["function"]["name"]}'
# assert len(tool_call.get("id", "")) > 0, f'Expected non empty tool call id in {tool_call}'
actual_arguments = json.loads(tool_call["function"]["arguments"])
assert 'location' in actual_arguments, f"location not found in {json.dumps(actual_arguments)}"
location = actual_arguments["location"]
assert isinstance(location, str), f"Expected location to be a string, got {type(location)}: {json.dumps(location)}"
assert re.match('^Istanbul(( |, ?)(TR|Turkey|Türkiye))?$', location), f'Expected Istanbul for location, got {location}'
@pytest.mark.slow
@pytest.mark.parametrize("stream", [CompletionMode.NORMAL, CompletionMode.STREAMED])
@pytest.mark.parametrize("result_override,n_predict,hf_repo,template_override", [
(None, 128, "bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", "chatml"),
(None, 128, "bartowski/Qwen2.5-Coder-3B-Instruct-GGUF:Q4_K_M", None),
(None, 128, "bartowski/Qwen2.5-Coder-3B-Instruct-GGUF:Q4_K_M", "chatml"),
(None, 128, "bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", "chatml"),
(None, 128, "bartowski/Hermes-2-Pro-Llama-3-8B-GGUF:Q4_K_M", ("NousResearch/Hermes-2-Pro-Llama-3-8B", "tool_use")),
(None, 128, "bartowski/Hermes-3-Llama-3.1-8B-GGUF:Q4_K_M", ("NousResearch/Hermes-3-Llama-3.1-8B", "tool_use")),
(None, 128, "bartowski/functionary-small-v3.2-GGUF:Q8_0", ("meetkai/functionary-medium-v3.2", None)),
(None, 128, "bartowski/Mistral-Nemo-Instruct-2407-GGUF:Q4_K_M", None),
(None, 128, "bartowski/Mistral-Nemo-Instruct-2407-GGUF:Q4_K_M", "chatml"),
(None, 128, "bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", None),
("[\\s\\S]*?\\*\\*\\s*0.5($|\\*\\*)", 8192, "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF:Q4_K_M", ("llama-cpp-deepseek-r1", None)),
# TODO: fix these (wrong results, either didn't respect decimal instruction or got wrong value)
# (None, 128, "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF:Q4_K_M", None),
# ("[\\s\\S]*?\\*\\*\\s*0.5($|\\*\\*)", 8192, "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF:Q4_K_M", None),
])
def test_calc_result(result_override: str | None, n_predict: int, hf_repo: str, template_override: str | Tuple[str, str | None] | None, stream: CompletionMode):
global server
server.jinja = True
server.n_ctx = 8192 * 2
server.n_predict = n_predict
server.model_hf_repo = hf_repo
server.model_hf_file = None
if isinstance(template_override, tuple):
(template_hf_repo, template_variant) = template_override
server.chat_template_file = f"../../../models/templates/{template_hf_repo.replace('/', '-') + ('-' + template_variant if template_variant else '')}.jinja"
assert os.path.exists(server.chat_template_file), f"Template file {server.chat_template_file} does not exist. Run `python scripts/get_chat_template.py {template_hf_repo} {template_variant} > {server.chat_template_file}` to download the template."
elif isinstance(template_override, str):
server.chat_template = template_override
server.start(timeout_seconds=TIMEOUT_START_SLOW)
do_test_calc_result(server, result_override, n_predict, stream=stream == CompletionMode.STREAMED)
def do_test_calc_result(server: ServerProcess, result_override: str | None, n_predict: int, **kwargs):
body = server.make_any_request("POST", "/v1/chat/completions", data={
"max_tokens": n_predict,
"messages": [
{"role": "system", "content": "You are a tools-calling assistant. You express numerical values with at most two decimals."},
{"role": "user", "content": "What's the y coordinate of a point on the unit sphere at angle 30 degrees?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_6789",
"type": "function",
"function": {
"name": "calculate",
"arguments": "{\"expression\":\"sin(30 * pi / 180)\"}"
}
}
]
},
{
"role": "tool",
"name": "calculate",
"content": "0.55644242476",
"tool_call_id": "call_6789"
}
],
"tools": [
{
"type":"function",
"function":{
"name":"calculate",
"description":"A calculator function that computes values of arithmetic expressions in the Python syntax",
"parameters":{
"type":"object",
"properties":{
"expression":{
"type":"string",
"description":"An arithmetic expression to compute the value of (Python syntad, assuming all floats)"
}
},
"required":["expression"]
}
}
}
],
**kwargs,
}, timeout=TIMEOUT_HTTP_REQUEST)
choice = body["choices"][0]
tool_calls = choice["message"].get("tool_calls")
assert tool_calls is None, f'Expected no tool call in {choice["message"]}'
content = choice["message"].get("content")
assert content is not None, f'Expected content in {choice["message"]}'
if result_override is not None:
assert re.match(result_override, content), f'Expected {result_override}, got {content}'
else:
assert re.match('^[\\s\\S]*?((That\'s|\\bis) (approximately )?)?\\b0\\.(5\\b|56\\b|556)', content), \
f'Expected something like "The y coordinate is 0.56.", got {content}'
@pytest.mark.slow
@pytest.mark.parametrize("stream", [CompletionMode.NORMAL, CompletionMode.STREAMED])
@pytest.mark.parametrize("n_predict,reasoning_format,expect_reasoning_content,expect_content,hf_repo,template_override", [
(128, 'deepseek', None, "^The sum of 102 and 7 is 109[\\s\\S]*", "bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", None),
(128, None, None, "^The sum of 102 and 7 is 109[\\s\\S]*", "bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", None),
(1024, 'deepseek', "I need to calculate the sum of 102 and 7[\\s\\S]*", "To find the sum of[\\s\\S]*", "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF:Q4_K_M", None),
(1024, 'deepseek', "First, I [\\s\\S]*", "To find the sum of[\\s\\S]*", "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF:Q4_K_M", ("llama-cpp-deepseek-r1", None)),
# (1024, 'none', CompletionMode.NORMAL, None, "^(<think>\\s*)?I need[\\s\\S]*?</think>\\s*To find[\\s\\S]*", "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF:Q4_K_M", None),
# (128, 'deepseek', None, "^Okay, let me figure out the sum of 102 and 7[\\s\\S]*", "bartowski/Qwen_QwQ-32B-GGUF:Q4_K_M", None),
])
def test_thoughts(n_predict: int, reasoning_format: Literal['deepseek', 'none'] | None, expect_content: str | None, expect_reasoning_content: str | None, hf_repo: str, template_override: str | Tuple[str, str | None] | None, stream: CompletionMode):
global server
server.reasoning_format = reasoning_format
server.jinja = True
server.n_ctx = 8192 * 2
server.n_predict = n_predict
server.model_hf_repo = hf_repo
server.model_hf_file = None
if isinstance(template_override, tuple):
(template_hf_repo, template_variant) = template_override
server.chat_template_file = f"../../../models/templates/{template_hf_repo.replace('/', '-') + ('-' + template_variant if template_variant else '')}.jinja"
assert os.path.exists(server.chat_template_file), f"Template file {server.chat_template_file} does not exist. Run `python scripts/get_chat_template.py {template_hf_repo} {template_variant} > {server.chat_template_file}` to download the template."
elif isinstance(template_override, str):
server.chat_template = template_override
server.start()
body = server.make_any_request("POST", "/v1/chat/completions", data={
"max_tokens": n_predict,
"messages": [
{"role": "user", "content": "What's the sum of 102 and 7?"},
],
"stream": stream == CompletionMode.STREAMED,
}, timeout=TIMEOUT_HTTP_REQUEST)
choice = body["choices"][0]
assert choice["message"].get("tool_calls") is None, f'Expected no tool call in {choice["message"]}'
content = choice["message"].get("content")
if expect_content is None:
assert choice["message"].get("content") in (None, ""), f'Expected no content in {choice["message"]}'
else:
assert re.match(expect_content, content), f'Expected {expect_content}, got {content}'
reasoning_content = choice["message"].get("reasoning_content")
if expect_reasoning_content is None:
assert reasoning_content is None, f'Expected no reasoning content in {choice["message"]}'
else:
assert re.match(expect_reasoning_content, reasoning_content), f'Expected {expect_reasoning_content}, got {reasoning_content}'
@pytest.mark.slow
@pytest.mark.parametrize("stream", [CompletionMode.NORMAL, CompletionMode.STREAMED])
@pytest.mark.parametrize("hf_repo,template_override", [
("bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF:Q4_K_M", None),
("bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", None),
("bartowski/Phi-3.5-mini-instruct-GGUF:Q4_K_M", "chatml"),
("bartowski/functionary-small-v3.2-GGUF:Q8_0", ("meetkai-functionary-medium-v3.2", None)),
("bartowski/functionary-small-v3.2-GGUF:Q8_0", "chatml"),
# ("bartowski/Meta-Llama-3.1-8B-Instruct-GGUF:Q4_K_M", None),
("bartowski/Meta-Llama-3.1-8B-Instruct-GGUF:Q4_K_M", "chatml"),
("bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M", ("meta-llama-Llama-3.2-3B-Instruct", None)),
("bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M", None),
("bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M", ("meta-llama-Llama-3.2-3B-Instruct", None)),
("bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M", None),
("bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", None),
("bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", "chatml"),
("bartowski/Hermes-2-Pro-Llama-3-8B-GGUF:Q4_K_M", ("NousResearch/Hermes-2-Pro-Llama-3-8B", "tool_use")),
("bartowski/Hermes-2-Pro-Llama-3-8B-GGUF:Q4_K_M", "chatml"),
("bartowski/Hermes-3-Llama-3.1-8B-GGUF:Q4_K_M", ("NousResearch-Hermes-3-Llama-3.1-8B", "tool_use")),
("bartowski/Hermes-3-Llama-3.1-8B-GGUF:Q4_K_M", "chatml"),
("bartowski/Mistral-Nemo-Instruct-2407-GGUF:Q4_K_M", None),
("bartowski/Mistral-Nemo-Instruct-2407-GGUF:Q4_K_M", "chatml"),
("bartowski/gemma-2-2b-it-GGUF:Q4_K_M", None),
("bartowski/gemma-2-2b-it-GGUF:Q4_K_M", "chatml"),
])
def test_hello_world(hf_repo: str, template_override: str | Tuple[str, str | None] | None, stream: CompletionMode):
global server
n_predict = 512 # High because of DeepSeek R1
server.jinja = True
server.n_ctx = 8192
server.n_predict = n_predict
server.model_hf_repo = hf_repo
server.model_hf_file = None
if isinstance(template_override, tuple):
(template_hf_repo, template_variant) = template_override
server.chat_template_file = f"../../../models/templates/{template_hf_repo.replace('/', '-') + ('-' + template_variant if template_variant else '')}.jinja"
assert os.path.exists(server.chat_template_file), f"Template file {server.chat_template_file} does not exist. Run `python scripts/get_chat_template.py {template_hf_repo} {template_variant} > {server.chat_template_file}` to download the template."
elif isinstance(template_override, str):
server.chat_template = template_override
server.start(timeout_seconds=TIMEOUT_START_SLOW)
do_test_hello_world(server, stream=stream == CompletionMode.STREAMED, max_tokens=n_predict)
def do_test_hello_world(server: ServerProcess, **kwargs):
body = server.make_any_request("POST", "/v1/chat/completions", data={
"messages": [
{"role": "system", "content": "You are a tool-calling agent."},
{"role": "user", "content": "say hello world with python"},
],
"tools": [PYTHON_TOOL],
**kwargs,
}, timeout=TIMEOUT_HTTP_REQUEST)
choice = body["choices"][0]
tool_calls = choice["message"].get("tool_calls")
assert tool_calls and len(tool_calls) == 1, f'Expected 1 tool call in {choice["message"]}'
tool_call = tool_calls[0]
# assert choice["message"].get("content") in (None, ""), f'Expected no content in {choice["message"]}'
assert tool_call["function"]["name"] == PYTHON_TOOL["function"]["name"]
# assert len(tool_call.get("id", "")) > 0, f'Expected non empty tool call id in {tool_call}'
actual_arguments = json.loads(tool_call["function"]["arguments"])
assert 'code' in actual_arguments, f"code not found in {json.dumps(actual_arguments)}"
code = actual_arguments["code"]
assert isinstance(code, str), f"Expected code to be a string, got {type(code)}: {json.dumps(code)}"
assert re.match(r'''print\(("[Hh]ello,? [Ww]orld!?"|'[Hh]ello,? [Ww]orld!?')\)''', re.sub(r'#.*\n?', '', code)), f'Expected hello world, got {code}'
+414
View File
@@ -0,0 +1,414 @@
import os
import shutil
import subprocess
import pytest
from utils import *
server: ServerProcess
# project root, used as the search directory for grep_search/file_glob_search
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))
# marker for the grep_search test to find in this file
GREP_MARKER = "llama_cpp_test_tools_builtin_marker_grep_search"
# image the container runtime tests run their shell in
CONTAINER_IMAGE = "busybox"
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.router()
server.server_tools = "all"
def call_tool(name: str, params: dict, headers: dict | None = None) -> dict:
res = server.make_request("POST", "/tools", data={"tool": name, "params": params}, headers=headers)
assert res.status_code == 200, res.body
assert "error" not in res.body, res.body
return res.body
def call_tool_expect_error(name: str, params: dict) -> str:
res = server.make_request("POST", "/tools", data={"tool": name, "params": params})
assert res.status_code == 200, res.body
assert "error" in res.body, res.body
return res.body["error"]
def test_tools_builtin_grep_search():
global server
server.start()
res = call_tool("grep_search", {
"path": PROJECT_ROOT,
"pattern": GREP_MARKER,
"include": "test_tools_builtin.py", # bare pattern -> matches basename at any depth
})
text = res["plain_text_response"]
assert "test_tools_builtin.py" in text
assert GREP_MARKER in text
assert "Total matches: 1" in text
def test_tools_builtin_read_file():
global server
server.start()
this_file = os.path.join(PROJECT_ROOT, "tools", "server", "tests", "unit", "test_tools_builtin.py")
res = call_tool("read_file", {"path": this_file})
text = res["plain_text_response"]
assert GREP_MARKER in text
assert "def test_tools_builtin_read_file" in text
def test_tools_builtin_write_then_edit_file(tmp_path):
global server
server.start()
log_path = str(tmp_path / "test.log")
try:
write_res = call_tool("write_file", {"path": log_path, "content": "line1\nline2\nline3\n"})
assert write_res["result"] == "file written successfully"
read_before = call_tool("read_file", {"path": log_path})
assert read_before["plain_text_response"] == "line1\nline2\nline3\n"
edit_res = call_tool("edit_file", {
"path": log_path,
"edits": [
{"old_text": "line2", "new_text": "line2-edited"},
{"old_text": "line3\n", "new_text": "line3\nline4\n"},
],
})
assert edit_res["result"] == "file edited successfully"
assert edit_res["edits_applied"] == 2
read_after = call_tool("read_file", {"path": log_path})
assert read_after["plain_text_response"] == "line1\nline2-edited\nline3\nline4\n"
finally:
if os.path.exists(log_path):
os.remove(log_path)
def test_tools_builtin_edit_file_rejects_non_unique_old_text(tmp_path):
global server
server.start()
log_path = str(tmp_path / "test.log")
try:
call_tool("write_file", {"path": log_path, "content": "dup\ndup\n"})
err = call_tool_expect_error("edit_file", {
"path": log_path,
"edits": [{"old_text": "dup", "new_text": "changed"}],
})
assert "unique" in err
finally:
if os.path.exists(log_path):
os.remove(log_path)
def test_tools_builtin_exec_shell_command_stream():
global server
server.start()
events = list(server.make_stream_request("POST", "/tools", data={
"tool": "exec_shell_command",
"params": {"command": "echo hello"},
"stream": True,
}))
assert len(events) >= 2
assert events[-1]["done"] is True
assert not events[-1].get("error")
chunks = "".join(e["chunk"] for e in events[:-1])
assert "hello" in chunks
assert "[exit code: 0]" in chunks
def test_tools_builtin_cwd_header():
global server
server.start()
cwd_dir = os.path.join(PROJECT_ROOT, "tools", "server", "tests", "unit")
headers = {"x-tool-cwd": cwd_dir}
res = call_tool("read_file", {"path": "test_tools_builtin.py"}, headers=headers)
assert GREP_MARKER in res["plain_text_response"]
# exec_shell_command should also run with that directory as its working directory:
# writing to a relative filename must land inside cwd_dir
marker_name = "llama_cpp_test_tools_builtin_cwd_marker.txt"
marker_path = os.path.join(cwd_dir, marker_name)
try:
command = f"echo hello > {marker_name}"
call_tool("exec_shell_command", {"command": command}, headers=headers)
assert os.path.exists(marker_path)
finally:
if os.path.exists(marker_path):
os.remove(marker_path)
def _container_engine_unavailable_reason(engine: str) -> str | None:
"""None if `engine` can run the image these tests use, otherwise the reason it can't."""
engine_bin = shutil.which(engine)
if engine_bin is None:
return f"{engine} is not installed"
try:
# a daemon that answers `info` still cannot run a linux image when it serves windows
# containers, so probe the image itself, which also pulls it before the tests
subprocess.run([engine_bin, "run", "--rm", CONTAINER_IMAGE, "true"], capture_output=True, timeout=60, check=True)
except Exception as e:
return f"{engine} cannot run {CONTAINER_IMAGE}: {e}"
return None
@pytest.fixture(params=["docker", "podman"])
def container_engine(request):
engine = request.param
reason = _container_engine_unavailable_reason(engine)
if reason is not None:
pytest.skip(reason) # ty: ignore[too-many-positional-arguments, invalid-argument-type]
return engine
@pytest.fixture
def container_id(container_engine: str):
proc = subprocess.run(
[container_engine, "run", "-d", "--rm", CONTAINER_IMAGE, "sleep", "300"],
capture_output=True, text=True,
)
if proc.returncode != 0:
pytest.skip(f"failed to start {container_engine} container: {proc.stderr.strip()}") # ty: ignore[too-many-positional-arguments, invalid-argument-type]
cid = proc.stdout.strip()
try:
yield cid
finally:
subprocess.run([container_engine, "rm", "-f", cid], capture_output=True)
def test_tools_builtin_runtime_header(container_engine: str, container_id: str):
global server
server.start()
headers = {"x-tool-runtime": f"{container_engine}-container:{container_id}", "x-tool-cwd": "/tmp"}
write_res = call_tool("write_file", {"path": "test.log", "content": "hello container\n"}, headers=headers)
assert write_res["result"] == "file written successfully"
read_res = call_tool("read_file", {"path": "test.log"}, headers=headers)
assert read_res["plain_text_response"] == "hello container\n"
exec_res = call_tool("exec_shell_command", {"command": "cat test.log"}, headers=headers)
assert "hello container" in exec_res["plain_text_response"]
def test_tools_builtin_runtime_header_unknown_scheme():
global server
server.start()
# an unknown runtime must fail, never silently fall back to running on the host
res = server.make_request("POST", "/tools",
data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},
headers={"x-tool-runtime": "fake:does-not-exist"})
assert res.status_code == 500, res.body
assert "unknown tool runtime" in str(res.body)
def test_tools_builtin_runtime_header_rejects_ssh_option_injection():
global server
server.start()
# ssh reads options from its argv, so a target starting with '-' must be rejected
res = server.make_request("POST", "/tools",
data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},
headers={"x-tool-runtime": "ssh:-oProxyCommand=touch /tmp/pwned"})
assert res.status_code == 500, res.body
assert "invalid ssh target" in str(res.body)
@pytest.mark.parametrize("engine", ["docker", "podman"])
def test_tools_builtin_runtime_header_rejects_container_option_injection(engine: str):
global server
server.start()
# the container id lands on the `<engine> exec` command line, so an id that looks
# like an option must be rejected
res = server.make_request("POST", "/tools",
data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},
headers={"x-tool-runtime": f"{engine}-container:--privileged"})
assert res.status_code == 500, res.body
assert "invalid container id" in str(res.body)
def test_tools_builtin_docker_runtime_cleans_up_spawned_container():
# docker-only: this reads the container hostname to get the spawned id, which only docker
# sets to the short id. podman is covered by the attach path above
reason = _container_engine_unavailable_reason("docker")
if reason is not None:
pytest.skip(reason) # ty: ignore[too-many-positional-arguments, invalid-argument-type]
global server
server.server_tools_runtime = f"docker:{CONTAINER_IMAGE}"
server.start()
# exec_shell_command runs inside the container spawned for --tools-runtime; docker sets
# the container's hostname to its own short id, so this also tells us which one to check
res = call_tool("exec_shell_command", {"command": "hostname"})
container_id = res["plain_text_response"].splitlines()[0].strip()
assert len(container_id) >= 8, res
running = subprocess.run(
["docker", "inspect", "-f", "{{.State.Running}}", container_id],
capture_output=True, text=True,
)
assert running.returncode == 0 and running.stdout.strip() == "true", running.stderr
server.stop()
# a clean server shutdown must stop and remove the container it spawned (it runs with --rm),
# not leave it behind as an abandoned child
leftover = subprocess.run(["docker", "inspect", container_id], capture_output=True, text=True)
assert leftover.returncode != 0, f"container {container_id} was not cleaned up after server exit"
def test_tools_builtin_edit_file_rejects_overlapping_edits(tmp_path):
global server
server.start()
log_path = str(tmp_path / "test.log")
try:
call_tool("write_file", {"path": log_path, "content": "line1\nline2\n"})
err = call_tool_expect_error("edit_file", {
"path": log_path,
"edits": [
{"old_text": "line1\nline2", "new_text": "a"},
{"old_text": "line2", "new_text": "b"},
],
})
assert "overlap" in err
finally:
if os.path.exists(log_path):
os.remove(log_path)
def test_tools_builtin_file_glob_search_type_dir(tmp_path):
global server
server.start()
(tmp_path / "project-alpha" / "src").mkdir(parents=True)
(tmp_path / "project-alpha" / "README.md").write_text("alpha")
(tmp_path / "project-alpha" / "src" / "main.cpp").write_text("int main() {}")
(tmp_path / "project-beta").mkdir()
(tmp_path / "project-beta" / "notes.txt").write_text("beta")
res = call_tool("file_glob_search", {"path": str(tmp_path), "type": "dir"})
text = res["plain_text_response"]
assert "project-alpha/" in text
assert "project-beta/" in text
assert "project-alpha/src/" in text
assert "README.md" not in text
types = {e["path"]: e["type"] for e in res["entries"]}
assert types["project-alpha"] == "dir"
assert types["project-alpha/src"] == "dir"
res_all = call_tool("file_glob_search", {"path": str(tmp_path), "type": "all", "include": "*proj*"})
paths = [e["path"] for e in res_all["entries"]]
assert "project-alpha" in paths
assert "project-beta" in paths
def test_tools_builtin_file_glob_search_max_depth_and_limit(tmp_path):
global server
server.start()
(tmp_path / "a" / "b" / "c").mkdir(parents=True)
(tmp_path / "top.txt").write_text("top")
(tmp_path / "a" / "mid.txt").write_text("mid")
(tmp_path / "a" / "b" / "deep.txt").write_text("deep")
res = call_tool("file_glob_search", {"path": str(tmp_path), "max_depth": 1})
assert "top.txt" in res["plain_text_response"]
assert "mid.txt" not in res["plain_text_response"]
res = call_tool("file_glob_search", {"path": str(tmp_path), "max_depth": 2})
assert "mid.txt" in res["plain_text_response"]
assert "deep.txt" not in res["plain_text_response"]
res = call_tool("file_glob_search", {"path": str(tmp_path), "limit": 1})
assert len(res["entries"]) == 1
assert "Total matches: 3" in res["plain_text_response"]
def test_tools_builtin_file_glob_search_junk_dirs(tmp_path):
global server
server.start()
(tmp_path / "build" / "nested").mkdir(parents=True)
(tmp_path / "build" / "artifact.txt").write_text("built")
(tmp_path / "src").mkdir()
(tmp_path / "src" / "main.cpp").write_text("int main() {}")
# a junk directory stays selectable as a working directory
res = call_tool("file_glob_search", {"path": str(tmp_path), "type": "dir", "max_depth": 1})
assert "build" in [e["path"] for e in res["entries"]]
# but it is never walked, so nothing inside it shows up
res = call_tool("file_glob_search", {"path": str(tmp_path), "type": "all"})
paths = [e["path"] for e in res["entries"]]
assert "src/main.cpp" in paths
assert "build/artifact.txt" not in paths
assert "build/nested" not in paths
def test_tools_builtin_file_glob_search_rejects_invalid_type(tmp_path):
global server
server.start()
err = call_tool_expect_error("file_glob_search", {"path": str(tmp_path), "type": "bogus"})
assert "invalid type" in err
def test_tools_builtin_cwd_header_overrides_model_param(tmp_path):
global server
server.start()
workdir = tmp_path / "workdir"
workdir.mkdir()
(workdir / "marker.txt").write_text("marker")
# a model-provided "cwd" in the params is overridden by the x-tool-cwd header
res = call_tool("read_file", {"path": "marker.txt", "cwd": "/definitely/not/a/real/path"},
headers={"x-tool-cwd": str(workdir)})
assert "marker" in res["plain_text_response"]
def test_tools_builtin_cwd_relative_paths(tmp_path):
global server
server.start()
workdir = tmp_path / "workdir"
workdir.mkdir()
(workdir / "rel.txt").write_text("relative-content")
headers = {"x-tool-cwd": str(workdir)}
# relative paths in file tools resolve against the header cwd
res = call_tool("read_file", {"path": "rel.txt"}, headers=headers)
assert "relative-content" in res["plain_text_response"]
res = call_tool("write_file", {"path": "sub/out.txt", "content": "written"}, headers=headers)
assert (workdir / "sub" / "out.txt").read_text() == "written"
res = call_tool("file_glob_search", {"path": ".", "include": "*.txt"}, headers=headers)
assert "rel.txt" in res["plain_text_response"]
# absolute paths are unaffected by the cwd
other = tmp_path / "other"
other.mkdir()
(other / "abs.txt").write_text("absolute-content")
res = call_tool("read_file", {"path": str(other / "abs.txt")}, headers=headers)
assert "absolute-content" in res["plain_text_response"]
@@ -0,0 +1,206 @@
import pytest
from utils import *
import base64
import requests
server: ServerProcess
def get_img_url(id: str) -> str:
IMG_URL_0 = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/11_truck.png"
IMG_URL_1 = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/91_cat.png"
if id == "IMG_URL_0":
return IMG_URL_0
elif id == "IMG_URL_1":
return IMG_URL_1
elif id == "IMG_BASE64_URI_0":
response = requests.get(IMG_URL_0)
response.raise_for_status() # Raise an exception for bad status codes
return "data:image/png;base64," + base64.b64encode(response.content).decode("utf-8")
elif id == "IMG_BASE64_0":
response = requests.get(IMG_URL_0)
response.raise_for_status() # Raise an exception for bad status codes
return base64.b64encode(response.content).decode("utf-8")
elif id == "IMG_BASE64_URI_1":
response = requests.get(IMG_URL_1)
response.raise_for_status() # Raise an exception for bad status codes
return "data:image/png;base64," + base64.b64encode(response.content).decode("utf-8")
elif id == "IMG_BASE64_1":
response = requests.get(IMG_URL_1)
response.raise_for_status() # Raise an exception for bad status codes
return base64.b64encode(response.content).decode("utf-8")
else:
return id
JSON_MULTIMODAL_KEY = "multimodal_data"
JSON_PROMPT_STRING_KEY = "prompt_string"
@pytest.fixture(autouse=True)
def create_server():
global server
os.environ['LLAMA_MEDIA_MARKER'] = '<__media__>'
server = ServerPreset.tinygemma3()
def test_models_supports_multimodal_capability():
global server
server.start()
res = server.make_request("GET", "/models", data={})
assert res.status_code == 200
model_info = res.body["models"][0]
print(model_info)
assert "completion" in model_info["capabilities"]
assert "multimodal" in model_info["capabilities"]
def test_v1_models_supports_multimodal_capability():
global server
server.start()
res = server.make_request("GET", "/v1/models", data={})
assert res.status_code == 200
model_info = res.body["models"][0]
print(model_info)
assert "completion" in model_info["capabilities"]
assert "multimodal" in model_info["capabilities"]
@pytest.mark.parametrize(
"prompt, image_url, success, re_content",
[
# test model is trained on CIFAR-10, but it's quite dumb due to small size
("What is this:\n", "IMG_URL_0", True, "(cat)+"),
("What is this:\n", "IMG_BASE64_URI_0", True, "(cat)+"),
("What is this:\n", "IMG_URL_1", True, "(frog)+"),
("Test test\n", "IMG_URL_1", True, "(frog)+"), # test invalidate cache
("What is this:\n", "malformed", False, None),
("What is this:\n", "https://google.com/404", False, None), # non-existent image
("What is this:\n", "https://ggml.ai", False, None), # non-image data
("What is this:\n", "data:text/html;base64,aGVsbG8=", False, None), # unsupported data uri mime
# TODO @ngxson : test with multiple images, no images and with audio
]
)
def test_vision_chat_completion(prompt, image_url, success, re_content):
global server
server.start()
res = server.make_request("POST", "/chat/completions", data={
"temperature": 0.0,
"top_k": 1,
"messages": [
{"role": "user", "content": [
{"type": "text", "text": prompt},
{"type": "image_url", "image_url": {
"url": get_img_url(image_url),
}},
]},
],
})
if success:
assert res.status_code == 200
choice = res.body["choices"][0]
assert "assistant" == choice["message"]["role"]
assert match_regex(re_content, choice["message"]["content"])
else:
assert res.status_code != 200
def test_vision_chat_completion_token_count():
global server
server.start()
res = server.make_request("POST", "/chat/completions/input_tokens", data={
"temperature": 0.0,
"top_k": 1,
"messages": [
{"role": "user", "content": [
{"type": "text", "text": "What is this:"},
{"type": "image_url", "image_url": {
"url": get_img_url("IMG_URL_0"),
}},
]},
],
})
assert res.status_code == 200
assert res.body["input_tokens"] > 10
@pytest.mark.parametrize(
"prompt, image_data, success, re_content",
[
# test model is trained on CIFAR-10, but it's quite dumb due to small size
("What is this: <__media__>\n", "IMG_BASE64_0", True, "(cat)+|(automobile)+"),
("What is this: <__media__>\n", "IMG_BASE64_1", True, "(frog)+"),
("What is this: <__media__>\n", "malformed", False, None), # non-image data
("What is this:\n", "", False, None), # empty string
]
)
def test_vision_completion(prompt, image_data, success, re_content):
global server
server.start()
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"prompt": {
JSON_PROMPT_STRING_KEY: prompt,
JSON_MULTIMODAL_KEY: [ get_img_url(image_data) ],
},
})
if success:
assert res.status_code == 200
content = res.body["content"]
assert match_regex(re_content, content)
else:
assert res.status_code != 200
@pytest.mark.parametrize(
"prompt, image_data, success",
[
# test model is trained on CIFAR-10, but it's quite dumb due to small size
("What is this: <__media__>\n", "IMG_BASE64_0", True),
("What is this: <__media__>\n", "IMG_BASE64_1", True),
("What is this: <__media__>\n", "malformed", False), # non-image data
("What is this:\n", "base64", False), # non-image data
]
)
def test_vision_embeddings(prompt, image_data, success):
global server
server.server_embeddings = True
server.n_batch = 512
server.start()
image_data = get_img_url(image_data)
res = server.make_request("POST", "/embeddings", data={
"content": [
{ JSON_PROMPT_STRING_KEY: prompt, JSON_MULTIMODAL_KEY: [ image_data ] },
{ JSON_PROMPT_STRING_KEY: prompt, JSON_MULTIMODAL_KEY: [ image_data ] },
{ JSON_PROMPT_STRING_KEY: prompt, },
],
})
if success:
assert res.status_code == 200
content = res.body
# Ensure embeddings are stable when multimodal.
assert content[0]['embedding'] == content[1]['embedding']
# Ensure embeddings without multimodal but same prompt do not match multimodal embeddings.
assert content[0]['embedding'] != content[2]['embedding']
else:
assert res.status_code != 200
def test_vision_embeddings_oai_content():
global server
server.server_embeddings = True
server.pooling = 'mean'
server.n_batch = 512
server.start()
res = server.make_request("POST", "/v1/embeddings", data={
"input": [
{"content": [
{"type": "text", "text": "What is this: "},
{"type": "image_url", "image_url": {"url": get_img_url("IMG_BASE64_URI_0")}},
{"type": "text", "text": "\n"},
]},
{JSON_PROMPT_STRING_KEY: "What is this: <__media__>\n", JSON_MULTIMODAL_KEY: [get_img_url("IMG_BASE64_0")]},
"What is this: \n",
],
})
assert res.status_code == 200
data = res.body["data"]
assert len(data) == 3
# same prompt and image in both formats
assert data[0]["embedding"] == data[1]["embedding"]
assert data[0]["embedding"] != data[2]["embedding"]
+778
View File
@@ -0,0 +1,778 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# type: ignore[reportUnusedImport]
import subprocess
import os
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
import re
import json
from json import JSONDecodeError
import sys
import requests
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import (
Any,
Callable,
ContextManager,
Iterable,
Iterator,
List,
Literal,
Tuple,
Set,
)
from re import RegexFlag
import wget
DEFAULT_HTTP_TIMEOUT = 600
# per-request timeout, a hung server fails the test instead of stalling the CI for hours
DEFAULT_REQUEST_TIMEOUT = 600
class ServerResponse:
headers: dict
status_code: int
body: dict | Any
class ServerError(Exception):
def __init__(self, code, body):
self.code = code
self.body = body
class ServerProcess:
# default options
debug: bool = False
server_port: int = 8080
server_host: str = "127.0.0.1"
model_hf_repo: str | None = "ggml-org/models"
model_hf_file: str | None = "tinyllamas/stories260K.gguf"
model_alias: str = "tinyllama-2"
temperature: float = 0.8
seed: int = 42
offline: bool = False
# custom options
model_alias: str | None = None
model_tags: str | None = None
model_url: str | None = None
model_file: str | None = None
model_draft: str | None = None
model_draft_hf_repo: str | None = None
n_threads: int | None = None
n_gpu_layer: int | None = None
n_batch: int | None = None
n_ubatch: int | None = None
n_ctx: int | None = None
n_ga: int | None = None
n_ga_w: int | None = None
n_predict: int | None = None
n_prompts: int | None = 0
slot_save_path: str | None = None
id_slot: int | None = None
cache_prompt: bool | None = None
n_slots: int | None = None
ctk: str | None = None
ctv: str | None = None
ctkd: str | None = None
fa: str | None = None
server_continuous_batching: bool | None = False
server_embeddings: bool | None = False
server_reranking: bool | None = False
server_metrics: bool | None = False
kv_unified: bool | None = False
swa_full: bool | None = False
server_slots: bool | None = False
pooling: str | None = None
api_key: str | None = None
models_dir: str | None = None
models_max: int | None = None
models_preset: str | None = None
no_models_autoload: bool | None = None
lora_files: List[str] | None = None
enable_ctx_shift: int | None = False
spec_type: str | None = None
spec_draft_n_min: int | None = None
spec_draft_n_max: int | None = None
spec_draft_n_ubatch: int | None = None
spec_synth_len: float | None = None
spec_synth_rates: List[float] | None = None
no_ui: bool | None = None
jinja: bool | None = None
reasoning_format: Literal['deepseek', 'none', 'nothink'] | None = None
reasoning: Literal['on', 'off', 'auto'] | None = None
chat_template: str | None = None
chat_template_file: str | None = None
server_path: str | None = None
mmproj_url: str | None = None
no_mmproj: bool | None = None
media_path: str | None = None
sleep_idle_seconds: int | None = None
cache_ram: int | None = None
override_kv: str | None = None
no_cache_idle_slots: bool = False
log_path: str | None = None
ui_mcp_proxy: bool = False
backend_sampling: bool = False
gcp_compat: bool = False
server_tools: str | None = None
server_tools_runtime: str | None = None
mcp_servers_config: str | None = None
mcp_servers_json: str | None = None
cors_origins: str | None = None
llama_cache: str | None = None
# session variables
process: subprocess.Popen | None = None
def __init__(self):
if "N_GPU_LAYERS" in os.environ:
self.n_gpu_layer = int(os.environ["N_GPU_LAYERS"])
if "DEBUG" in os.environ:
self.debug = True
if "PORT" in os.environ:
self.server_port = int(os.environ["PORT"])
self.external_server = "DEBUG_EXTERNAL" in os.environ
def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None:
env = {
**os.environ,
"LLAMA_SERVER_DEBUG_FAKE_TIMING": "1",
}
if self.llama_cache is not None:
env["LLAMA_CACHE"] = self.llama_cache
elif "LLAMA_CACHE" not in os.environ:
env["LLAMA_CACHE"] = "tmp"
if self.external_server:
print(f"[external_server]: Assuming external server running on {self.server_host}:{self.server_port}")
return
if self.server_path is not None:
server_path = self.server_path
elif "LLAMA_SERVER_BIN_PATH" in os.environ:
server_path = os.environ["LLAMA_SERVER_BIN_PATH"]
elif os.name == "nt":
server_path = "../../../build/bin/Release/llama-server.exe"
else:
server_path = "../../../build/bin/llama-server"
server_args = [
"--port",
self.server_port,
"--temp",
self.temperature,
"--seed",
self.seed,
]
server_args.extend(["--host", self.server_host])
if self.offline:
server_args.append("--offline")
if self.model_file:
server_args.extend(["--model", self.model_file])
if self.model_url:
server_args.extend(["--model-url", self.model_url])
if self.model_draft:
server_args.extend(["--model-draft", self.model_draft])
if self.model_draft_hf_repo:
server_args.extend(["--hf-repo-draft", self.model_draft_hf_repo])
if self.model_hf_repo:
server_args.extend(["--hf-repo", self.model_hf_repo])
if self.model_hf_file:
server_args.extend(["--hf-file", self.model_hf_file])
if self.models_dir:
server_args.extend(["--models-dir", self.models_dir])
if self.models_max is not None:
server_args.extend(["--models-max", self.models_max])
if self.models_preset:
server_args.extend(["--models-preset", self.models_preset])
if self.cors_origins:
server_args.extend(["--cors-origins", self.cors_origins])
if self.n_batch:
server_args.extend(["--batch-size", self.n_batch])
if self.n_ubatch:
server_args.extend(["--ubatch-size", self.n_ubatch])
if self.n_threads:
server_args.extend(["--threads", self.n_threads])
if self.n_gpu_layer:
server_args.extend(["--n-gpu-layers", self.n_gpu_layer])
if self.server_continuous_batching:
server_args.append("--cont-batching")
if self.server_embeddings:
server_args.append("--embedding")
if self.server_reranking:
server_args.append("--reranking")
if self.server_metrics:
server_args.append("--metrics")
if self.kv_unified:
server_args.append("--kv-unified")
if self.swa_full:
server_args.append("--swa-full")
if self.server_slots:
server_args.append("--slots")
else:
server_args.append("--no-slots")
if self.pooling:
server_args.extend(["--pooling", self.pooling])
if self.model_alias:
server_args.extend(["--alias", self.model_alias])
if self.model_tags:
server_args.extend(["--tags", self.model_tags])
if self.n_ctx:
server_args.extend(["--ctx-size", self.n_ctx])
if self.n_slots:
server_args.extend(["--parallel", self.n_slots])
if self.ctk:
server_args.extend(["-ctk", self.ctk])
if self.ctv:
server_args.extend(["-ctv", self.ctv])
if self.ctkd:
server_args.extend(["-ctkd", self.ctkd])
if self.fa is not None:
server_args.extend(["-fa", self.fa])
if self.n_predict:
server_args.extend(["--n-predict", self.n_predict])
if self.slot_save_path:
server_args.extend(["--slot-save-path", self.slot_save_path])
if self.n_ga:
server_args.extend(["--grp-attn-n", self.n_ga])
if self.n_ga_w:
server_args.extend(["--grp-attn-w", self.n_ga_w])
if self.debug:
server_args.append("--verbose")
if self.lora_files:
for lora_file in self.lora_files:
server_args.extend(["--lora", lora_file])
if self.enable_ctx_shift:
server_args.append("--context-shift")
if self.spec_type:
server_args.extend(["--spec-type", self.spec_type])
if self.api_key:
server_args.extend(["--api-key", self.api_key])
if self.spec_draft_n_max:
server_args.extend(["--spec-draft-n-max", self.spec_draft_n_max])
if self.spec_draft_n_min:
server_args.extend(["--spec-draft-n-min", self.spec_draft_n_min])
if self.spec_draft_n_ubatch:
server_args.extend(["--spec-draft-n-ubatch", self.spec_draft_n_ubatch])
if self.spec_synth_len is not None:
server_args.extend(["--spec-synth-len", self.spec_synth_len])
if self.spec_synth_rates is not None:
rates = ",".join(str(rate) for rate in self.spec_synth_rates)
server_args.extend(["--spec-synth-rates", rates])
if self.no_ui:
server_args.append("--no-ui")
if self.no_models_autoload:
server_args.append("--no-models-autoload")
if self.jinja:
server_args.append("--jinja")
else:
server_args.append("--no-jinja")
if self.reasoning_format is not None:
server_args.extend(("--reasoning-format", self.reasoning_format))
if self.reasoning is not None:
server_args.extend(("--reasoning", self.reasoning))
if self.chat_template:
server_args.extend(["--chat-template", self.chat_template])
if self.chat_template_file:
server_args.extend(["--chat-template-file", self.chat_template_file])
if self.mmproj_url:
server_args.extend(["--mmproj-url", self.mmproj_url])
if self.no_mmproj:
server_args.append("--no-mmproj")
if self.media_path:
server_args.extend(["--media-path", self.media_path])
if self.sleep_idle_seconds is not None:
server_args.extend(["--sleep-idle-seconds", self.sleep_idle_seconds])
if self.cache_ram is not None:
server_args.extend(["--cache-ram", self.cache_ram])
if self.override_kv:
server_args.extend(["--override-kv", self.override_kv])
if self.no_cache_idle_slots:
server_args.append("--no-cache-idle-slots")
if self.ui_mcp_proxy:
server_args.append("--ui-mcp-proxy")
if self.server_tools:
server_args.extend(["--tools", self.server_tools])
if self.server_tools_runtime:
server_args.extend(["--tools-runtime", self.server_tools_runtime])
if self.mcp_servers_config:
server_args.extend(["--mcp-servers-config", self.mcp_servers_config])
if self.mcp_servers_json:
server_args.extend(["--mcp-servers-json", self.mcp_servers_json])
if self.backend_sampling:
server_args.append("--backend_sampling")
if self.gcp_compat:
env["AIP_MODE"] = "PREDICTION"
env["AIP_HTTP_PORT"] = str(self.server_port)
args = [str(arg) for arg in [server_path, *server_args]]
print(f"tests: starting server with: {' '.join(args)}")
flags = 0
if "nt" == os.name:
flags |= subprocess.DETACHED_PROCESS
flags |= subprocess.CREATE_NEW_PROCESS_GROUP
flags |= subprocess.CREATE_NO_WINDOW
if self.log_path:
self._log = open(self.log_path, "w")
else:
self._log = sys.stdout
self.process = subprocess.Popen(
[str(arg) for arg in [server_path, *server_args]],
creationflags=flags,
stdout=self._log,
stderr=self._log if self._log != sys.stdout else sys.stdout,
env=env,
)
server_instances.add(self)
print(f"server pid={self.process.pid}, pytest pid={os.getpid()}")
# wait for server to start
start_time = time.time()
last_print_time = start_time
while time.time() - start_time < timeout_seconds:
try:
response = self.make_request("GET", "/health", headers={
"Authorization": f"Bearer {self.api_key}" if self.api_key else None
})
if response.status_code == 200:
self.ready = True
return # server is ready
except Exception as e:
pass
# Check if process died
if self.process.poll() is not None:
raise RuntimeError(f"Server process died with return code {self.process.returncode}")
if time.time() - last_print_time >= 1.0:
print(f"Waiting for server to start...")
last_print_time = time.time()
time.sleep(0.01)
raise TimeoutError(f"Server did not start within {timeout_seconds} seconds")
def stop(self) -> None:
if self.external_server:
print("[external_server]: Not stopping external server")
return
if self in server_instances:
server_instances.remove(self)
if self.process:
print(f"Stopping server with pid={self.process.pid}")
self.process.terminate()
try:
self.process.wait(timeout=5)
except subprocess.TimeoutExpired:
print(f"Server pid={self.process.pid} did not terminate in time, killing")
self.process.kill()
self.process.wait(timeout=5)
except Exception as e:
print(f"Error waiting for server: {e}")
self.process = None
if hasattr(self, '_log') and self._log != sys.stdout:
self._log.close()
def make_url(self, path: str, host: str | None = None) -> str:
if host is None:
host = self.server_host.split(",")[0].strip()
return f"http://{host}:{self.server_port}{path}"
def make_request(
self,
method: str,
path: str,
data: dict | Any | None = None,
headers: dict | None = None,
timeout: float | None = DEFAULT_REQUEST_TIMEOUT,
host: str | None = None,
) -> ServerResponse:
url = self.make_url(path, host)
parse_body = False
if method == "GET":
response = requests.get(url, headers=headers, timeout=timeout)
parse_body = True
elif method == "POST":
response = requests.post(url, headers=headers, json=data, timeout=timeout)
parse_body = True
elif method == "DELETE":
response = requests.delete(url, headers=headers, timeout=timeout)
parse_body = True
elif method == "OPTIONS":
response = requests.options(url, headers=headers, timeout=timeout)
else:
raise ValueError(f"Unimplemented method: {method}")
result = ServerResponse()
result.headers = dict(response.headers)
result.status_code = response.status_code
if parse_body:
try:
result.body = response.json()
except (JSONDecodeError, requests.exceptions.JSONDecodeError):
result.body = response.text
else:
result.body = None
print("Response from server", json.dumps(result.body, indent=2))
return result
def make_stream_request(
self,
method: str,
path: str,
data: dict | None = None,
headers: dict | None = None,
host: str | None = None,
) -> Iterator[dict]:
url = self.make_url(path, host)
if method == "POST":
response = requests.post(url, headers=headers, json=data, stream=True)
else:
raise ValueError(f"Unimplemented method: {method}")
if response.status_code != 200:
raise ServerError(response.status_code, response.json())
for line_bytes in response.iter_lines():
line = line_bytes.decode("utf-8")
if '[DONE]' in line:
break
elif line.startswith('data: '):
data = json.loads(line[6:])
print("Partial response from server", json.dumps(data, indent=2))
yield data
def make_any_request(
self,
method: str,
path: str,
data: dict | None = None,
headers: dict | None = None,
timeout: float | None = DEFAULT_REQUEST_TIMEOUT,
) -> dict:
stream = data.get('stream', False)
if stream:
content: list[str] = []
reasoning_content: list[str] = []
tool_calls: list[dict] = []
finish_reason: Optional[str] = None
content_parts = 0
reasoning_content_parts = 0
tool_call_parts = 0
arguments_parts = 0
for chunk in self.make_stream_request(method, path, data, headers):
if chunk['choices']:
assert len(chunk['choices']) == 1, f'Expected 1 choice, got {len(chunk["choices"])}'
choice = chunk['choices'][0]
if choice['delta'].get('content') is not None:
assert len(choice['delta']['content']) > 0, f'Expected non empty content delta!'
content.append(choice['delta']['content'])
content_parts += 1
if choice['delta'].get('reasoning_content') is not None:
assert len(choice['delta']['reasoning_content']) > 0, f'Expected non empty reasoning_content delta!'
reasoning_content.append(choice['delta']['reasoning_content'])
reasoning_content_parts += 1
if choice['delta'].get('finish_reason') is not None:
finish_reason = choice['delta']['finish_reason']
for tc in choice['delta'].get('tool_calls', []):
if 'function' not in tc:
raise ValueError(f"Expected function type, got {tc['type']}")
if tc['index'] >= len(tool_calls):
assert 'id' in tc
assert tc.get('type') == 'function'
assert 'function' in tc and 'name' in tc['function'] and len(tc['function']['name']) > 0, \
f"Expected function call with name, got {tc.get('function')}"
tool_calls.append(dict(
id="",
type="function",
function=dict(
name="",
arguments="",
)
))
tool_call = tool_calls[tc['index']]
if tc.get('id') is not None:
tool_call['id'] = tc['id']
fct = tc['function']
assert 'id' not in fct, f"Function call should not have id: {fct}"
if fct.get('name') is not None:
tool_call['function']['name'] = tool_call['function'].get('name', '') + fct['name']
if fct.get('arguments') is not None:
tool_call['function']['arguments'] += fct['arguments']
arguments_parts += 1
tool_call_parts += 1
else:
# When `include_usage` is True (the default), we expect the last chunk of the stream
# immediately preceding the `data: [DONE]` message to contain a `choices` field with an empty array
# and a `usage` field containing the usage statistics (n.b., llama-server also returns `timings` in
# the last chunk)
assert 'usage' in chunk, f"Expected finish_reason in chunk: {chunk}"
assert 'timings' in chunk, f"Expected finish_reason in chunk: {chunk}"
print(f'Streamed response had {content_parts} content parts, {reasoning_content_parts} reasoning_content parts, {tool_call_parts} tool call parts incl. {arguments_parts} arguments parts')
result = dict(
choices=[
dict(
index=0,
finish_reason=finish_reason,
message=dict(
role='assistant',
content=''.join(content) if content else None,
reasoning_content=''.join(reasoning_content) if reasoning_content else None,
tool_calls=tool_calls if tool_calls else None,
),
)
],
)
print("Final response from server", json.dumps(result, indent=2))
return result
else:
response = self.make_request(method, path, data, headers, timeout=timeout)
assert response.status_code == 200, f"Server returned error: {response.status_code}"
return response.body
server_instances: Set[ServerProcess] = set()
class ServerPreset:
@staticmethod
def load_all() -> None:
""" Load all server presets to ensure model files are cached. """
servers: List[ServerProcess] = [
method()
for name, method in ServerPreset.__dict__.items()
if callable(method) and name != "load_all"
]
for server in servers:
server.offline = False
server.start()
server.stop()
@staticmethod
def tinyllama2() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
server.model_hf_repo = "ggml-org/test-model-stories260K"
server.model_hf_file = None
server.model_alias = "tinyllama-2"
server.n_ctx = 512
server.n_batch = 32
server.n_slots = 2
server.n_predict = 64
server.seed = 42
return server
@staticmethod
def bert_bge_small() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
server.model_hf_repo = "ggml-org/models"
server.model_hf_file = "bert-bge-small/ggml-model-f16.gguf"
server.model_alias = "bert-bge-small"
server.n_ctx = 512
server.n_batch = 128
server.n_ubatch = 128
server.n_slots = 2
server.seed = 42
server.server_embeddings = True
return server
@staticmethod
def bert_bge_small_with_fa() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
server.model_hf_repo = "ggml-org/models"
server.model_hf_file = "bert-bge-small/ggml-model-f16.gguf"
server.model_alias = "bert-bge-small"
server.n_ctx = 1024
server.n_batch = 300
server.n_ubatch = 300
server.n_slots = 2
server.fa = "on"
server.seed = 42
server.server_embeddings = True
return server
@staticmethod
def tinyllama_infill() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
server.model_hf_repo = "ggml-org/test-model-stories260K-infill"
server.model_hf_file = None
server.model_alias = "tinyllama-infill"
server.n_ctx = 2048
server.n_batch = 1024
server.n_slots = 1
server.n_predict = 64
server.temperature = 0.0
server.seed = 42
return server
@staticmethod
def stories15m_moe() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
server.model_hf_repo = "ggml-org/stories15M_MOE"
server.model_hf_file = "stories15M_MOE-F16.gguf"
server.model_alias = "stories15m-moe"
server.n_ctx = 2048
server.n_batch = 1024
server.n_slots = 1
server.n_predict = 64
server.temperature = 0.0
server.seed = 42
return server
@staticmethod
def jina_reranker_tiny() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
server.model_hf_repo = "ggml-org/models"
server.model_hf_file = "jina-reranker-v1-tiny-en/ggml-model-f16.gguf"
server.model_alias = "jina-reranker"
server.n_ctx = 512
server.n_batch = 512
server.n_slots = 1
server.seed = 42
server.server_reranking = True
return server
@staticmethod
def tinylaya() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
local_model = os.environ.get("TINYLAYA_LOCAL_MODEL")
server.model_hf_file = None
if local_model:
server.model_file = local_model
server.model_hf_repo = None
else:
server.model_hf_repo = "ggml-org/tinylaya-for-testing-gguf"
server.n_ctx = 1024
server.n_batch = 512
server.n_ubatch = 512
server.n_slots = 2
server.seed = 42
return server
@staticmethod
def tinyopenjev() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
# mmproj is already provided by HF registry API
server.model_hf_file = None
server.model_hf_repo = "ggml-org/tinyopenjev-for-testing-gguf:Q8_0"
server.n_ctx = 4096
server.n_batch = 512
server.n_slots = 4
server.seed = 42
return server
@staticmethod
def tinygemma3() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
# mmproj is already provided by HF registry API
server.model_hf_file = None
server.model_hf_repo = "ggml-org/tinygemma3-GGUF:Q8_0"
server.model_alias = "tinygemma3"
server.n_ctx = 1024
server.n_batch = 512
server.n_slots = 2
server.n_predict = 4
server.seed = 42
return server
@staticmethod
def router() -> ServerProcess:
server = ServerProcess()
server.offline = True # will be downloaded by load_all()
# router server has no models
server.model_file = None
server.model_alias = None
server.model_hf_repo = None
server.model_hf_file = None
server.n_ctx = 1024
server.n_batch = 16
server.n_slots = 1
server.n_predict = 16
server.seed = 42
return server
def parallel_function_calls(function_list: List[Tuple[Callable[..., Any], Tuple[Any, ...]]]) -> List[Any]:
"""
Run multiple functions in parallel and return results in the same order as calls. Equivalent to Promise.all in JS.
Example usage:
results = parallel_function_calls([
(func1, (arg1, arg2)),
(func2, (arg3, arg4)),
])
"""
results = [None] * len(function_list)
exceptions = []
def worker(index, func, args):
try:
result = func(*args)
results[index] = result
except Exception as e:
exceptions.append((index, str(e)))
with ThreadPoolExecutor() as executor:
futures = []
for i, (func, args) in enumerate(function_list):
future = executor.submit(worker, i, func, args)
futures.append(future)
# Wait for all futures to complete
for future in as_completed(futures):
pass
# Check if there were any exceptions
if exceptions:
print("Exceptions occurred:")
for index, error in exceptions:
print(f"Function at index {index}: {error}")
return results
def match_regex(regex: str, text: str) -> bool:
return (
re.compile(
regex, flags=RegexFlag.IGNORECASE | RegexFlag.MULTILINE | RegexFlag.DOTALL
).search(text)
is not None
)
def download_file(url: str, output_file_path: str | None = None) -> str:
"""
Download a file from a URL to a local path. If the file already exists, it will not be downloaded again.
output_file_path is the local path to save the downloaded file. If not provided, the file will be saved in the root directory.
Returns the local path of the downloaded file.
"""
file_name = url.split('/').pop()
output_file = f'./tmp/{file_name}' if output_file_path is None else output_file_path
if not os.path.exists(output_file):
print(f"Downloading {url} to {output_file}")
wget.download(url, out=output_file)
print(f"Done downloading to {output_file}")
else:
print(f"File already exists at {output_file}")
return output_file
def is_slow_test_allowed():
return os.environ.get("SLOW_TESTS") == "1" or os.environ.get("SLOW_TESTS") == "ON"