other.prompt = prompt.clone();
other.init_sampler();
}
+};
- // returns 0 on success
- // caller need to update prompt.tokens after a successful call to keep track of the processing progress
- int process_mtmd_chunk(size_t idx, size_t & n_tokens_out) {
- GGML_ASSERT(mctx);
- const auto & input_tokens = task->tokens;
- const auto & chunk = input_tokens.find_chunk(idx);
- int32_t res = 0;
-
- auto try_decode = [&]() -> int32_t {
- if (mbatch) {
- float * embd = mtmd_batch_get_output_embd(mbatch.get(), chunk.get());
- if (embd) {
- void * cb_data = spec;
- static auto cb = [](llama_batch batch, void * user_data) {
- common_speculative * spec = static_cast<common_speculative *>(user_data);
- if (!common_speculative_process(spec, batch)) {
- return 1;
- }
- return 0;
- };
-
- llama_pos new_n_past; // unused for now
- res = mtmd_helper_decode_image_chunk(
- mctx,
- ctx_tgt,
- chunk.get(),
- embd,
- prompt.tokens.pos_next(),
- id,
- llama_n_batch(ctx_tgt),
- &new_n_past,
- cb,
- cb_data
- );
- if (res != 0) {
- SLT_ERR(*this, "failed to decode mtmd chunk, idx = %zu, res = %d\n", idx, res);
- return -1;
+// returns 0 on success
+// caller need to update prompt.tokens after a successful call to keep track of the processing progress
+// note: this is not a member of server_slot because we want to run it inside yield_to_queue
+// slot is passed as const to avoid accidental modification of the slot state
+// some pointers are allowed to be used, they are not used by to_json()
+static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch, size_t idx, size_t & n_tokens_out) {
+ GGML_ASSERT(slot.mctx);
+ const auto & mctx = slot.mctx;
+ const auto & input_tokens = slot.task->tokens;
+ const auto & chunk = input_tokens.find_chunk(idx);
+ int32_t res = 0;
+
+ auto try_decode = [&]() -> int32_t {
+ if (mbatch) {
+ float * embd = mtmd_batch_get_output_embd(mbatch.get(), chunk.get());
+ if (embd) {
+ void * cb_data = slot.spec;
+ static auto cb = [](llama_batch batch, void * user_data) {
+ common_speculative * spec = static_cast<common_speculative *>(user_data);
+ if (!common_speculative_process(spec, batch)) {
+ return 1;
}
- n_tokens_out = mtmd_input_chunk_get_n_tokens(chunk.get());
- return 0; // success
+ return 0;
+ };
+
+ llama_pos new_n_past; // unused for now
+ res = mtmd_helper_decode_image_chunk(
+ mctx,
+ slot.ctx_tgt,
+ chunk.get(),
+ embd,
+ slot.prompt.tokens.pos_next(),
+ slot.id,
+ llama_n_batch(slot.ctx_tgt),
+ &new_n_past,
+ cb,
+ cb_data
+ );
+ if (res != 0) {
+ SLT_ERR(slot, "failed to decode mtmd chunk, idx = %zu, res = %d\n", idx, res);
+ return -1;
}
+ n_tokens_out = mtmd_input_chunk_get_n_tokens(chunk.get());
+ return 0; // success
}
- return 1; // (non-error) need to create & encode batch
- };
-
- // if the batch is already exist, try searching & encode
- res = try_decode();
- if (res == 0) {
- return 0;
- }
- if (res < 0) {
- // fatal error
- return res;
}
+ return 1; // (non-error) need to create & encode batch
+ };
- // otherwise, the batch is either uninitialized or is used up
- // we need to create & encode a new batch
- mbatch.reset(mtmd_batch_init(mctx));
- res = mtmd_batch_add_chunk(mbatch.get(), chunk.get());
- GGML_ASSERT(res == 0); // we should never have an empty batch
-
- // try batching as much as possible
- int n_added = 1;
- size_t idx_cur = idx;
- while (res == 0) {
- auto [next_chunk, next_idx] = input_tokens.find_next_media_chunk(idx_cur);
- if (next_chunk == nullptr) {
- break;
- }
- res = mtmd_batch_add_chunk(mbatch.get(), next_chunk->get());
- n_added += (res == 0 ? 1 : 0);
- idx_cur = next_idx;
- SLT_DBG(*this, "try adding media chunk idx = %zu to batch, res = %d\n", next_idx, res);
- // if res != 0, batch is full or chunk is not compatible -> this loop breaks
- }
+ // if the batch is already exist, try searching & encode
+ res = try_decode();
+ if (res == 0) {
+ return 0;
+ }
+ if (res < 0) {
+ // fatal error
+ return res;
+ }
- // TODO @ngxson : move this log line to debug when it become more stable
- SLT_TRC(*this, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added);
+ // otherwise, the batch is either uninitialized or is used up
+ // we need to create & encode a new batch
+ mbatch.reset(mtmd_batch_init(mctx));
+ res = mtmd_batch_add_chunk(mbatch.get(), chunk.get());
+ GGML_ASSERT(res == 0); // we should never have an empty batch
+
+ // try batching as much as possible
+ int n_added = 1;
+ size_t idx_cur = idx;
+ while (res == 0) {
+ auto [next_chunk, next_idx] = input_tokens.find_next_media_chunk(idx_cur);
+ if (next_chunk == nullptr) {
+ break;
+ }
+ res = mtmd_batch_add_chunk(mbatch.get(), next_chunk->get());
+ n_added += (res == 0 ? 1 : 0);
+ idx_cur = next_idx;
+ SLT_DBG(slot, "try adding media chunk idx = %zu to batch, res = %d\n", next_idx, res);
+ // if res != 0, batch is full or chunk is not compatible -> this loop breaks
+ }
- res = mtmd_batch_encode(mbatch.get());
- if (res != 0) {
- SLT_ERR(*this, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res);
- return -1;
- }
+ // TODO @ngxson : move this log line to debug when it become more stable
+ SLT_TRC(slot, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added);
- return try_decode();
+ res = mtmd_batch_encode(mbatch.get());
+ if (res != 0) {
+ SLT_ERR(slot, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res);
+ return -1;
}
-};
-
+ return try_decode();
+}
//
// server_context_impl (private implementation)
GGML_ASSERT(!sleeping);
// wiring up server queues
- queue_tasks.on_new_task([this](server_task && task) {
- process_single_task(std::move(task));
+ queue_tasks.on_new_task([this](server_task && task, bool is_yielding) {
+ return process_single_task(std::move(task), is_yielding);
});
queue_tasks.on_update_slots([this]() {
update_slots();
cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
}
- void process_single_task(server_task && task) {
+ // returns false to decline the task, it is offered again after the decode is done
+ bool process_single_task(server_task && task, bool is_yielding) {
+ // while yielding, an encode / decode is running and only accessing metrics is safe
+ if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS) {
+ SRV_DBG("decoding, decline task, id_task = %d\n", task.id);
+ return false;
+ }
+
switch (task.type) {
case SERVER_TASK_TYPE_COMPLETION:
case SERVER_TASK_TYPE_INFILL:
queue_results.send(std::move(res));
} break;
}
+
+ return true;
}
void iterate(std::vector<server_slot> & slots, std::function<void(server_slot &)> callback) {
// so the timing is queued and flushed on the next sync
metrics_pre_decode();
+ // encode on the worker thread, so we can still handle metrics tasks
size_t n_tokens_out = 0;
- int32_t res = slot.process_mtmd_chunk(cur_token_idx, n_tokens_out);
+ int32_t res = 0;
+ queue_tasks.yield_to_queue([&]() {
+ res = process_mtmd_chunk(slot, slot.mbatch, cur_token_idx, n_tokens_out);
+ });
+
if (res != 0) {
SLT_ERR(slot, "failed to process mtmd chunk, res = %d\n", res);
send_error(slot, "failed to process mtmd chunk", ERROR_TYPE_SERVER);
}
}
- const int ret = llama_decode(ctx_tgt, batch_view);
+ bool has_output = false;
+ for (int i = off; i < off + batch_view.n_tokens; ++i) {
+ has_output |= batch.tokens[i].output;
+ }
+
+ // decode on the worker thread, so we can still handle metrics tasks while waiting
+ // note: the sync is done here too, so that the wait also happens off the main thread
+ int ret = 0;
+ queue_tasks.yield_to_queue([&]() {
+ ret = llama_decode(ctx_tgt, batch_view);
+ if (ret == 0 && has_output) {
+ llama_synchronize(ctx_tgt);
+ }
+ });
if (ret != 0) {
{
return false; // retry with the updated n_batch
} else {
// success, apply batch metrics
- metrics_post_decode(off, batch_view.n_tokens);
+ metrics_post_decode(off, batch_view.n_tokens, has_output);
}
// TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL]
n_prompt_queued = 0;
}
- void metrics_post_decode(int32_t off, int32_t n_tokens) {
+ // has_output is computed by the caller, which also already synchronized the context if it is set
+ void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) {
metrics.n_decode++;
for (const auto & slot : slots) {
if (slot.is_processing()) {
// note: a slot can be released before we get here, which clears its stats
// the tokens were still computed, counted in the global metrics, not in slot
uint64_t n_prompt_tokens = 0;
- bool has_output = false;
for (int i = off; i < off + n_tokens; ++i) {
const auto & t = batch.tokens[i];
- has_output |= t.output;
-
if (!t.is_prompt) {
continue; // generated tokens are handled after sampling
}
metrics_queue_prompt(n_prompt_tokens);
if (has_output) {
- // sync if we have at least one output in batch
- // so that we can calculate the timings correctly
- llama_synchronize(ctx_tgt);
+ // the context is already synchronized, so the timings are correct
metrics_flush_prompt();
}
// advance the prompt timing of the slots that had tokens in this batch
- // note: a second pass, it must run after the sync above to reflect the compute
+ // note: a second pass, it must run after the sync to reflect the compute
const int64_t t_now = ggml_time_us();
for (int i = off; i < off + n_tokens; ++i) {
const auto & t = batch.tokens[i];
#include "log.h"
#include <chrono>
+#include <thread>
#define QUE_INF(fmt, ...) LOG_INF("que %12.*s: " fmt, 12, __func__, __VA_ARGS__)
#define QUE_WRN(fmt, ...) LOG_WRN("que %12.*s: " fmt, 12, __func__, __VA_ARGS__)
condition_tasks.notify_all();
}
+bool server_queue::process_new_tasks(bool is_yielding) {
+ while (true) {
+ std::unique_lock<std::mutex> lock(mutex_tasks);
+ if (!running) {
+ QUE_DBG("%s", "terminate\n");
+ return true;
+ }
+ if (queue_tasks.empty()) {
+ return false;
+ }
+ server_task task = std::move(queue_tasks.front());
+ queue_tasks.pop_front();
+ lock.unlock();
+
+ QUE_DBG("processing task, id = %d\n", task.id);
+ if (!callback_new_task(std::move(task), is_yielding)) {
+ // set it aside, do not put it back in the queue, else we offer it again in a loop
+ GGML_ASSERT(is_yielding && "a task can only be declined while yielding");
+ QUE_DBG("task declined, id = %d\n", task.id);
+ lock.lock();
+ queue_tasks_unhandled.push_back(std::move(task));
+ }
+ }
+}
+
+void server_queue::worker_loop() {
+ while (true) {
+ std::function<void()> work;
+ {
+ std::unique_lock<std::mutex> lock(mutex_tasks);
+ worker.cv.wait(lock, [&]{
+ return worker.stop || worker.work != nullptr;
+ });
+ if (worker.stop) {
+ return;
+ }
+ work = std::move(worker.work);
+ worker.work = nullptr;
+ }
+
+ // note: do not hold any lock here, work() may post new tasks
+ std::exception_ptr exception;
+ try {
+ work();
+ } catch (...) {
+ exception = std::current_exception();
+ }
+
+ // signal completion to yield_to_queue()
+ std::unique_lock<std::mutex> lock(mutex_tasks);
+ worker.exception = std::move(exception);
+ worker.busy = false;
+ condition_tasks.notify_all();
+ }
+}
+
+void server_queue::worker_stop() {
+ if (!worker.thread.joinable()) {
+ return;
+ }
+ {
+ std::unique_lock<std::mutex> lock(mutex_tasks);
+ worker.stop = true;
+ }
+ worker.cv.notify_one();
+ worker.thread.join();
+}
+
+void server_queue::yield_to_queue(std::function<void()> && work) {
+ GGML_ASSERT(worker.thread.joinable() && "yield_to_queue() requires start_loop() to be running");
+
+ QUE_DBG("%s", "yielding to queue\n");
+
+ {
+ std::unique_lock<std::mutex> lock(mutex_tasks);
+ GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested");
+ worker.busy = true;
+ worker.work = std::move(work);
+ }
+ worker.cv.notify_one();
+
+ while (true) {
+ // note: on terminate this is a no-op, but we still wait for the work to finish
+ process_new_tasks(true);
+
+ std::unique_lock<std::mutex> lock(mutex_tasks);
+ // declined tasks are moved to queue_tasks_unhandled, so a non-empty queue always has something new
+ condition_tasks.wait(lock, [&]{
+ return !worker.busy || (running && !queue_tasks.empty());
+ });
+ if (!worker.busy) {
+ break;
+ }
+ }
+
+ std::exception_ptr exception;
+ {
+ std::unique_lock<std::mutex> lock(mutex_tasks);
+
+ // put the declined tasks back, keeping their order
+ while (!queue_tasks_unhandled.empty()) {
+ queue_tasks.push_front(std::move(queue_tasks_unhandled.back()));
+ queue_tasks_unhandled.pop_back();
+ }
+
+ // make sure to avoid idle timeout here
+ time_last_task = ggml_time_ms();
+
+ // the worker is idle now, take the exception it may have left behind
+ std::swap(exception, worker.exception);
+ }
+
+ QUE_DBG("%s", "done yielding to queue\n");
+
+ // note: rethrow only after the declined tasks are back in the queue, so they are not lost
+ if (exception) {
+ std::rethrow_exception(exception);
+ }
+}
+
void server_queue::start_loop(int64_t idle_sleep_ms) {
running = true;
time_last_task = ggml_time_ms();
+ // spawn the worker thread used by yield_to_queue()
+ GGML_ASSERT(!worker.thread.joinable() && "start_loop() is already running");
+ worker.stop = false;
+ worker.thread = std::thread([this]() { worker_loop(); });
+
constexpr auto max_wait_time = std::chrono::seconds(1);
auto should_sleep = [&]() -> bool {
// caller must hold mutex_tasks
while (true) {
QUE_DBG("%s", "processing new tasks\n");
-
- while (true) {
- std::unique_lock<std::mutex> lock(mutex_tasks);
- if (!running) {
- QUE_DBG("%s", "terminate\n");
- return;
- }
- if (queue_tasks.empty()) {
- lock.unlock();
- break;
- }
- server_task task = std::move(queue_tasks.front());
- queue_tasks.pop_front();
- lock.unlock();
-
- QUE_DBG("processing task, id = %d\n", task.id);
- callback_new_task(std::move(task));
+ if (process_new_tasks(false)) {
+ break; // terminate
}
+
// all tasks in the current loop is processed, slots data is now ready
QUE_DBG("%s", "update slots\n");
}
}
}
+
+ worker_stop();
}
void server_queue::cleanup_pending_task(int id_target) {
return task.id == id_target;
};
queue_tasks.erase(
- std::remove_if(queue_tasks.begin(), queue_tasks.end(), rm_func),
+ std::remove_if(queue_tasks.begin(), queue_tasks.end(), rm_func),
queue_tasks.end());
queue_tasks_deferred.erase(
- std::remove_if(queue_tasks_deferred.begin(), queue_tasks_deferred.end(), rm_func),
+ std::remove_if(queue_tasks_deferred.begin(), queue_tasks_deferred.end(), rm_func),
queue_tasks_deferred.end());
+ // a task declined while yielding is not in queue_tasks yet, but it can still be cancelled
+ queue_tasks_unhandled.erase(
+ std::remove_if(queue_tasks_unhandled.begin(), queue_tasks_unhandled.end(), rm_func),
+ queue_tasks_unhandled.end());
}
//
#include <condition_variable>
#include <deque>
+#include <exception>
#include <mutex>
+#include <thread>
#include <vector>
#include <unordered_set>
// queues
std::deque<server_task> queue_tasks;
std::deque<server_task> queue_tasks_deferred;
+ // tasks declined while yielding, put back in queue_tasks once the yield is done
+ // note: kept as a member so that cleanup_pending_task() can also reach them
+ std::deque<server_task> queue_tasks_unhandled;
std::mutex mutex_tasks;
std::condition_variable condition_tasks;
+ // used by yield_to_queue, all fields are guarded by mutex_tasks
+ struct worker_t {
+ std::thread thread;
+ std::condition_variable cv; // the worker sleeps on this until there is work
+ std::function<void()> work; // pending work, picked up by the thread
+ std::exception_ptr exception; // exception thrown by work(), if any
+ bool stop = false;
+ bool busy = false;
+ };
+ worker_t worker;
+
// callback functions
- std::function<void(server_task &&)> callback_new_task;
- std::function<void(void)> callback_update_slots;
- std::function<void(bool)> callback_sleeping_state;
+ std::function<bool(server_task &&, bool)> callback_new_task;
+ std::function<void(void)> callback_update_slots;
+ std::function<void(bool)> callback_sleeping_state;
public:
+ ~server_queue() { worker_stop(); }
+
// Add a new task to the end of the queue
int post(server_task && task, bool front = false);
*/
void start_loop(int64_t idle_sleep_ms = -1);
+ // run work() on a separate thread, while the current thread calls process_new_tasks
+ // returns once work() is done (may throw exceptions)
+ // must be called from start_loop() thread (ideally inside callback_update_slots)
+ // use case: return metrics while encode/decode is running
+ // ref: https://github.com/ggml-org/llama.cpp/pull/27041
+ //
+ // tasks declined by callback_new_task are put back in the queue once this returns
+ void yield_to_queue(std::function<void()> && work);
+
// for metrics
size_t queue_tasks_deferred_size() {
std::unique_lock<std::mutex> lock(mutex_tasks);
//
// Register function to process a new task
- void on_new_task(std::function<void(server_task &&)> callback) {
+ // the second argument tells whether the queue is currently yielding (see yield_to_queue)
+ // only then may the callback return false to decline the task, and it must leave it
+ // untouched, so that it can be put back in the queue later
+ void on_new_task(std::function<bool(server_task &&, bool)> callback) {
callback_new_task = std::move(callback);
}
private:
void cleanup_pending_task(int id_target);
+
+ // process all pending tasks in the queue
+ // returns true if the queue is terminated, false if there is no more task to process
+ // while yielding, declined tasks are moved to queue_tasks_unhandled
+ bool process_new_tasks(bool is_yielding);
+
+ // for worker_t
+ void worker_loop();
+ void worker_stop();
};
// struct for managing server responses
overview / pkg / ggml / sources / llama.cpp / commitdiff