From 77918caf30e65f830ab0759c89779efdac28c77b Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Fri, 14 Aug 2026 13:23:10 +0200 Subject: [PATCH] server: allow accessing /metrics and /slots during llama_decode() (#27041) * server_queue::worker * call llama_decode inside yield_to_queue * also handle process_mtmd_chunk * clean up * nits * rm test --- tools/server/server-context.cpp | 229 ++++++++++++++++++-------------- tools/server/server-queue.cpp | 156 +++++++++++++++++++--- tools/server/server-queue.h | 47 ++++++- 3 files changed, 307 insertions(+), 125 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f02a1da68..a7d382886 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -688,97 +688,99 @@ struct server_slot { 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(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; - } - 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; - } - - // 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 - } - - // 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); - - 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; - } - - return try_decode(); - } }; +// 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(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, + 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; + } + + // 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 + } + + // 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); + + 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) @@ -1354,8 +1356,8 @@ private: 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(); @@ -2286,7 +2288,14 @@ private: 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: @@ -2620,6 +2629,8 @@ private: queue_results.send(std::move(res)); } break; } + + return true; } void iterate(std::vector & slots, std::function callback) { @@ -3382,8 +3393,13 @@ private: // 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); @@ -3557,7 +3573,20 @@ private: } } - 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) { { @@ -3609,7 +3638,7 @@ private: 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] @@ -3922,7 +3951,8 @@ private: 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()) { @@ -3935,13 +3965,10 @@ private: // 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 } @@ -3957,14 +3984,12 @@ private: 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]; diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 5d37c3453..d22527ffe 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -4,6 +4,7 @@ #include "log.h" #include +#include #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__) @@ -122,10 +123,135 @@ void server_queue::terminate() { condition_tasks.notify_all(); } +bool server_queue::process_new_tasks(bool is_yielding) { + while (true) { + std::unique_lock 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 work; + { + std::unique_lock 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 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 lock(mutex_tasks); + worker.stop = true; + } + worker.cv.notify_one(); + worker.thread.join(); +} + +void server_queue::yield_to_queue(std::function && 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 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 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 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 @@ -138,24 +264,10 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { while (true) { QUE_DBG("%s", "processing new tasks\n"); - - while (true) { - std::unique_lock 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"); @@ -206,6 +318,8 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { } } } + + worker_stop(); } void server_queue::cleanup_pending_task(int id_target) { @@ -214,11 +328,15 @@ 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()); } // diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index 0b674d6ff..69a3e554f 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -4,7 +4,9 @@ #include #include +#include #include +#include #include #include @@ -21,16 +23,32 @@ private: // queues std::deque queue_tasks; std::deque 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 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 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 callback_new_task; - std::function callback_update_slots; - std::function callback_sleeping_state; + std::function callback_new_task; + std::function callback_update_slots; + std::function 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); @@ -75,6 +93,15 @@ public: */ 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 && work); + // for metrics size_t queue_tasks_deferred_size() { std::unique_lock lock(mutex_tasks); @@ -86,7 +113,10 @@ public: // // Register function to process a new task - void on_new_task(std::function 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 callback) { callback_new_task = std::move(callback); } @@ -112,6 +142,15 @@ public: 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