Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions tools/server/server-common.h
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,9 @@ struct server_tokens {

bool empty() const { return tokens.empty(); }

// true if the sequence actually contains image/audio chunks.
bool has_media() const { return !map_idx_to_media.empty(); }

void clear() {
map_idx_to_media.clear();
tokens.clear();
Expand Down
30 changes: 16 additions & 14 deletions tools/server/server-context.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2056,10 +2056,13 @@ struct server_context_impl {
queue_results.send(std::move(res));
}

// if multimodal is enabled, send an error and return false
bool check_no_mtmd(const int id_task) {
if (mctx) {
send_error(id_task, "This feature is not supported by multimodal", ERROR_TYPE_NOT_SUPPORTED);
// Gate slot save/restore/erase on slot content (does it hold media),
// not model capability: a multimodal model may hold a pure-text slot.
bool check_slot_no_media(const server_slot & slot, const int id_task) {
if (slot.prompt.tokens.has_media()) {
send_error(id_task,
"This operation is not supported while the slot holds image/audio tokens (a pure-text prefix is supported)",
ERROR_TYPE_NOT_SUPPORTED);
return false;
}
return true;
Expand Down Expand Up @@ -2525,30 +2528,29 @@ struct server_context_impl {
} break;
case SERVER_TASK_TYPE_SLOT_SAVE:
{
if (!check_no_mtmd(task.id)) {
break;
}

const int id_slot = task.slot_action.id_slot;
server_slot * slot = get_slot_by_id(id_slot);
if (slot == nullptr) {
send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST);
break;
}
if (!check_slot_no_media(*slot, task.id)) {
break;
}
if (slot->is_processing()) {
// if requested slot is unavailable, we defer this task for processing later
SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id);
queue_tasks.defer(std::move(task));
break;
}

const size_t token_count = slot->prompt.tokens.size();
const int64_t t_start = ggml_time_us();

std::string filename = task.slot_action.filename;
std::string filepath = task.slot_action.filepath;

const llama_tokens & tokens = slot->prompt.tokens.get_tokens();
const llama_tokens tokens = slot->prompt.tokens.get_text_tokens();
const size_t token_count = tokens.size();
const size_t nwrite = llama_state_seq_save_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), token_count);

const int64_t t_end = ggml_time_us();
Expand All @@ -2566,7 +2568,6 @@ struct server_context_impl {
} break;
case SERVER_TASK_TYPE_SLOT_RESTORE:
{
if (!check_no_mtmd(task.id)) break;
const int id_slot = task.slot_action.id_slot;
server_slot * slot = get_slot_by_id(id_slot);
if (slot == nullptr) {
Expand Down Expand Up @@ -2613,15 +2614,16 @@ struct server_context_impl {
} break;
case SERVER_TASK_TYPE_SLOT_ERASE:
{
if (!check_no_mtmd(task.id)) {
break;
}
const int id_slot = task.slot_action.id_slot;
server_slot * slot = get_slot_by_id(id_slot);
if (slot == nullptr) {
send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST);
break;
}
// Gate on slot content, consistent with save/restore.
if (!check_slot_no_media(*slot, task.id)) {
break;
}
if (slot->is_processing()) {
// if requested slot is unavailable, we defer this task for processing later
SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id);
Expand Down
126 changes: 126 additions & 0 deletions tools/server/tests/unit/test_slot_save.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import pytest
from utils import *
import base64
import requests

server = ServerPreset.tinyllama2()

Expand Down Expand Up @@ -96,3 +98,127 @@ def test_slot_erase():
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.
#
# Regression coverage for issue #21133: slot save/restore/erase must be gated on
# the slot's CONTENT (does it actually hold image/audio tokens) rather than the
# model's CAPABILITY (is an mmproj loaded). A pure-text slot on a multimodal
# server must save/restore/erase normally; a slot that actually holds an image
# must be rejected with ERROR_TYPE_NOT_SUPPORTED (HTTP 501).
#

IMG_URL_CAT = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/91_cat.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

# The restored slot is usable for a follow-up completion. We do NOT assert
# prefix reuse here: tinygemma3 is a SWA model, which forces full prompt
# re-processing after a restore (a model property, not the save/restore gate
# under test).
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_rejected_when_slot_holds_image(mmproj_server):
server = mmproj_server
server.start()

# Process a prompt that actually contains an image on slot 1.
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 1,
"cache_prompt": True,
"prompt": {
"prompt_string": "What is this: <__media__>\n",
"multimodal_data": [ _get_img_base64(IMG_URL_CAT) ],
},
})
assert res.status_code == 200

# Saving a slot that holds image tokens must be rejected (HTTP 501,
# not_supported_error).
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": "mm_slot_image.bin",
})
assert res.status_code != 200
assert res.body["error"]["type"] == "not_supported_error"


def test_slot_erase_text_only_on_multimodal(mmproj_server):
server = mmproj_server
server.start()

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

# Erasing a pure-text slot must succeed even though an mmproj is loaded.
res = server.make_request("POST", "/slots/1?action=erase")
assert res.status_code == 200

# Re-running the same prompt should process all tokens again.
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
assert res.body["timings"]["prompt_n"] == prompt_n # all tokens are processed again
Loading