#include <sstream>
#include <fstream>
#include <limits>
+#include <cstring>
+#include <type_traits>
json format_error_response(const std::string & message, const enum error_type type) {
std::string type_str;
// server_tokens implementation
//
+namespace {
+
+constexpr uint32_t SERVER_TOKENS_STATE_VERSION = 1;
+
+uint32_t server_tokens_state_u32(size_t value) {
+ if (value > std::numeric_limits<uint32_t>::max()) {
+ throw std::runtime_error("Server tokens state is too large");
+ }
+ return value;
+}
+
+class server_tokens_state_writer {
+public:
+ template <typename T>
+ void write(T value) {
+ static_assert(std::is_trivially_copyable<T>::value, "T must be trivially copyable");
+ const auto * ptr = reinterpret_cast<const char *>(&value);
+ data.insert(data.end(), ptr, ptr + sizeof(value));
+ }
+
+ template <typename T>
+ void write(const std::vector<T> & values) {
+ static_assert(std::is_trivially_copyable<T>::value, "T must be trivially copyable");
+ write(server_tokens_state_u32(values.size()));
+ if (values.empty()) {
+ return;
+ }
+ const auto * ptr = reinterpret_cast<const char *>(values.data());
+ data.insert(data.end(), ptr, ptr + values.size() * sizeof(T));
+ }
+
+ void write_media_chunk(const mtmd_input_chunk * chunk) {
+ size_t chunk_size = 0;
+ if (mtmd_input_chunk_save(chunk, nullptr, 0, &chunk_size) != 0 || chunk_size == 0) {
+ throw std::runtime_error("Cannot serialize media chunk in server tokens");
+ }
+ std::vector<char> chunk_data(server_tokens_state_u32(chunk_size));
+ if (mtmd_input_chunk_save(chunk, chunk_data.data(), chunk_data.size(), nullptr) != 0) {
+ throw std::runtime_error("Cannot serialize media chunk in server tokens");
+ }
+ write(chunk_data);
+ }
+
+ std::vector<char> take() {
+ data.resize((data.size() + sizeof(llama_token) - 1) / sizeof(llama_token) * sizeof(llama_token), 0);
+ return std::move(data);
+ }
+
+private:
+ std::vector<char> data;
+};
+
+class server_tokens_state_reader {
+public:
+ server_tokens_state_reader(const char * data, size_t size) : data(data), size(size) {}
+
+ template <typename T>
+ T read() {
+ static_assert(std::is_trivially_copyable<T>::value, "T must be trivially copyable");
+ if (size - pos < sizeof(T)) {
+ throw std::runtime_error("Unexpected end of server tokens state");
+ }
+ T value;
+ std::memcpy(&value, data + pos, sizeof(value));
+ pos += sizeof(value);
+ return value;
+ }
+
+ template <typename T>
+ std::vector<T> read_vector() {
+ static_assert(std::is_trivially_copyable<T>::value, "T must be trivially copyable");
+ const uint32_t n_values = read<uint32_t>();
+ // reject before resizing, so that a small corrupted payload cannot request a huge allocation
+ if (n_values > remaining() / sizeof(T)) {
+ throw std::runtime_error("Unexpected end of server tokens state");
+ }
+ std::vector<T> values(n_values);
+ if (n_values > 0) {
+ std::memcpy(values.data(), data + pos, values.size() * sizeof(T));
+ pos += values.size() * sizeof(T);
+ }
+ return values;
+ }
+
+ size_t remaining() const {
+ return size - pos;
+ }
+
+private:
+ const char * data;
+ size_t size;
+ size_t pos = 0;
+};
+
+} // namespace
+
server_tokens::server_tokens(mtmd::input_chunks & mtmd_chunks, bool has_mtmd) : has_mtmd(has_mtmd) {
for (size_t i = 0; i < mtmd_chunks.size(); ++i) {
push_back(mtmd_chunks[i]);
return tokens;
}
+std::vector<char> server_tokens::serialize() const {
+ static_assert(sizeof(llama_token) == sizeof(uint32_t), "unexpected llama_token size");
+
+ server_tokens_state_writer writer;
+ writer.write((llama_token) LLAMA_TOKEN_NULL);
+ writer.write(SERVER_TOKENS_STATE_VERSION);
+ writer.write(tokens);
+
+ std::vector<uint32_t> media_keys;
+ media_keys.reserve(map_idx_to_media.size());
+ for (const auto & item : map_idx_to_media) {
+ media_keys.push_back(server_tokens_state_u32(item.first));
+ }
+ writer.write(media_keys);
+
+ for (const auto & item : map_idx_to_media) {
+ writer.write_media_chunk(item.second.get());
+ }
+
+ return writer.take();
+}
+
+server_tokens server_tokens::deserialize(const llama_tokens & packed, bool has_mtmd) {
+ static_assert(sizeof(llama_token) == sizeof(uint32_t), "unexpected llama_token size");
+
+ if (packed.empty() || packed[0] != LLAMA_TOKEN_NULL) {
+ // plain token list, as written by older versions
+ return server_tokens(packed, has_mtmd);
+ }
+
+ server_tokens_state_reader reader(reinterpret_cast<const char *>(packed.data()), packed.size() * sizeof(llama_token));
+ reader.read<llama_token>(); // format marker
+ if (reader.read<uint32_t>() != SERVER_TOKENS_STATE_VERSION) {
+ throw std::runtime_error("Unsupported server tokens state version");
+ }
+
+ const llama_tokens tokens = reader.read_vector<llama_token>();
+
+ // the media start indices, followed by the media chunks in the same order
+ const std::vector<uint32_t> media_keys = reader.read_vector<uint32_t>();
+ if (!media_keys.empty() && !has_mtmd) {
+ throw std::runtime_error("Cannot restore media tokens without an mmproj");
+ }
+
+ server_tokens result(tokens, has_mtmd);
+
+ for (const uint32_t key : media_keys) {
+ const size_t start_idx = key;
+ const std::vector<char> chunk_data = reader.read_vector<char>();
+ if (chunk_data.empty()) {
+ throw std::runtime_error("Cannot load media chunk from server tokens state");
+ }
+
+ mtmd::input_chunk_ptr chunk(mtmd_input_chunk_load(chunk_data.data(), chunk_data.size()));
+ if (!chunk) {
+ throw std::runtime_error("Cannot load media chunk from server tokens state");
+ }
+ result.map_idx_to_media[start_idx] = std::move(chunk);
+ }
+
+ if (reader.remaining() >= sizeof(llama_token)) {
+ throw std::runtime_error("Trailing data in server tokens state");
+ }
+
+ return result;
+}
+
llama_tokens server_tokens::get_text_tokens() const {
llama_tokens res;
res.reserve(tokens.size());
const llama_model * model = llama_get_model(ctx);
const llama_vocab * vocab = llama_model_get_vocab(model);
const int32_t n_vocab = llama_vocab_n_tokens(vocab);
+ size_t n_media = 0;
for (size_t i = 0; i < tokens.size(); ++i) {
const auto & t = tokens[i];
if (t == LLAMA_TOKEN_NULL) {
try {
const auto & chunk = find_chunk(i);
- size_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk.get());
- i += n_tokens - 1; // will be +1 by the for loop
+ if (mtmd_input_chunk_get_type(chunk.get()) == MTMD_INPUT_CHUNK_TYPE_TEXT) {
+ return false;
+ }
+ const size_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk.get());
+ const llama_pos n_pos = mtmd_input_chunk_get_n_pos(chunk.get());
+ if (n_tokens == 0 || n_pos <= 0 || n_tokens > tokens.size() - i) {
+ return false;
+ }
+ for (size_t j = i; j < i + n_tokens; ++j) {
+ if (tokens[j] != LLAMA_TOKEN_NULL) {
+ return false;
+ }
+ }
+ ++n_media;
+ i += n_tokens - 1;
} catch (const std::exception & e) {
return false;
}
return false;
}
}
- return true;
+ return n_media == map_idx_to_media.size();
}
server_tokens server_tokens::clone() const {
queue_results.send(std::move(res));
}
- // 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;
- }
-
void send_partial_response(server_slot & slot, const completion_token_output & tkn, bool is_progress, bool is_begin = false) {
auto res = std::make_unique<server_task_result_cmpl_partial>();
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);
std::string filename = task.slot_action.filename;
std::string filepath = task.slot_action.filepath;
- 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);
+ std::vector<char> packed;
+ try {
+ packed = slot->prompt.tokens.serialize();
+ } catch (const std::exception & err) {
+ send_error(task, err.what(), ERROR_TYPE_NOT_SUPPORTED);
+ break;
+ }
+
+ GGML_ASSERT(packed.size() % sizeof(llama_token) == 0);
+ const size_t nwrite = llama_state_seq_save_file(
+ ctx_tgt, filepath.c_str(), slot->id,
+ reinterpret_cast<const llama_token *>(packed.data()), packed.size() / sizeof(llama_token));
+ if (nwrite == 0) {
+ send_error(task, "Unable to save slot", ERROR_TYPE_SERVER);
+ break;
+ }
const int64_t t_end = ggml_time_us();
const double t_save_ms = (t_end - t_start) / 1000.0;
res->id_slot = id_slot;
res->filename = filename;
res->is_save = true;
- res->n_tokens = token_count;
+ res->n_tokens = slot->prompt.tokens.size();
res->n_bytes = nwrite;
res->t_ms = t_save_ms;
queue_results.send(std::move(res));
std::string filename = task.slot_action.filename;
std::string filepath = task.slot_action.filepath;
- llama_tokens tokens;
- tokens.resize(slot->n_ctx);
- size_t token_count = 0;
- size_t nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), tokens.size(), &token_count);
- if (nread == 0) {
- slot->prompt.clear(); // KV may already been invalidated?
- send_error(task, "Unable to restore slot, no available space in KV cache or invalid slot save file", ERROR_TYPE_INVALID_REQUEST);
+ size_t nread = 0;
+ try {
+ size_t n_packed = 0;
+ llama_tokens packed;
+ nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, nullptr, 0, &n_packed);
+ if (nread != 0) {
+ packed.resize(std::max<size_t>(1, n_packed));
+ nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, packed.data(), packed.size(), &n_packed);
+ }
+ if (nread == 0) {
+ throw std::runtime_error("No available space in KV cache or invalid slot save file");
+ }
+ packed.resize(n_packed);
+
+ server_tokens restored = server_tokens::deserialize(packed, mctx != nullptr);
+
+ if (restored.size() > (size_t) slot->n_ctx) {
+ throw std::runtime_error("Restored prompt does not fit in the slot context");
+ }
+
+ if (!restored.validate(ctx_tgt)) {
+ throw std::runtime_error("Invalid tokens in slot save file");
+ }
+
+ slot->prompt.clear();
+ slot->prompt.tokens = std::move(restored);
+ } catch (const std::exception & err) {
+ slot->prompt_clear();
+ send_error(task, std::string("Unable to restore slot: ") + err.what(), ERROR_TYPE_INVALID_REQUEST);
break;
}
- tokens.resize(token_count);
- slot->prompt.clear();
- slot->prompt.tokens.insert(tokens);
const int64_t t_end = ggml_time_us();
const double t_restore_ms = (t_end - t_start) / 1000.0;
res->id_slot = id_slot;
res->filename = filename;
res->is_save = false;
- res->n_tokens = token_count;
+ res->n_tokens = slot->prompt.tokens.size();
res->n_bytes = nread;
res->t_ms = t_restore_ms;
queue_results.send(std::move(res));
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);
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()
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("tmp", "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
+
+
+
def test_slot_erase():
global server
server.start()
#
# 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).
+# 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:
@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.
+ # 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"
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).
+ # 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,
assert res.status_code == 200
-def test_slot_save_rejected_when_slot_holds_image(mmproj_server):
+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()
- # Process a prompt that actually contains an image on slot 1.
+ 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_CAT) ],
+ "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
- # 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={
+ # 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["error"]["type"] == "not_supported_error"
+ 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_erase_text_only_on_multimodal(mmproj_server):
+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()
- res = server.make_request("POST", "/completion", data={
- "prompt": "The quick brown fox jumps over the lazy dog.",
+ 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 = res.body["timings"]["prompt_n"]
- assert prompt_n > 0 # all tokens are processed
+ content = res.body["content"]
+ prompt_n_full = res.body["timings"]["prompt_n"]
+ assert prompt_n_full > 64
- # Erasing a pure-text slot must succeed even though an mmproj is loaded.
- res = server.make_request("POST", "/slots/1?action=erase")
+ 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
+ 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("tmp", "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
- # Re-running the same prompt should process all tokens again.
+ # drop the image from the slot, then restore it from the file
res = server.make_request("POST", "/completion", data={
- "prompt": "The quick brown fox jumps over the lazy dog.",
+ "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"]["prompt_n"] == prompt_n # all tokens are processed again
+ assert res.body["timings"]["cache_n"] == 0
+ assert res.body["content"] == content