diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f02a1da687..a45ae7a456 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -854,6 +854,9 @@ private: int slots_debug = 0; int n_empty_consecutive = 0; + // true while a decode runs on the yield_to_queue() worker thread + bool is_decoding = false; + std::unique_ptr prompt_cache; server_metrics metrics; @@ -1355,7 +1358,7 @@ private: // wiring up server queues queue_tasks.on_new_task([this](server_task && task) { - process_single_task(std::move(task)); + return process_single_task(std::move(task)); }); queue_tasks.on_update_slots([this]() { update_slots(); @@ -2286,7 +2289,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) { + // only metrics is safe while decoding, it touches neither ctx_tgt nor the slots + if (is_decoding && 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 +2630,8 @@ private: queue_results.send(std::move(res)); } break; } + + return true; } void iterate(std::vector & slots, std::function callback) { @@ -3557,7 +3569,21 @@ 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 + int ret = 0; + is_decoding = true; + queue_tasks.yield_to_queue([&]() { + ret = llama_decode(ctx_tgt, batch_view); + if (ret == 0 && has_output) { + llama_synchronize(ctx_tgt); + } + }); + is_decoding = false; if (ret != 0) { { diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 2826319faa..0c82a73f99 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -123,7 +123,7 @@ void server_queue::terminate() { condition_tasks.notify_all(); } -bool server_queue::process_new_tasks() { +bool server_queue::process_new_tasks(std::deque * unhandled) { while (true) { std::unique_lock lock(mutex_tasks); if (!running) { @@ -138,7 +138,12 @@ bool server_queue::process_new_tasks() { lock.unlock(); QUE_DBG("processing task, id = %d\n", task.id); - callback_new_task(std::move(task)); + if (!callback_new_task(std::move(task))) { + // set it aside, do not put it back in the queue, else we offer it again in a loop + GGML_ASSERT(unhandled && "a task can only be declined while yielding"); + QUE_DBG("task declined, id = %d\n", task.id); + unhandled->push_back(std::move(task)); + } } } @@ -164,7 +169,7 @@ void server_queue::worker_loop() { worker.exception = std::current_exception(); } - // signal completion to the thread waiting in yield_to_queue() + // signal completion to yield_to_queue() std::unique_lock lock(mutex_tasks); worker.busy = false; condition_tasks.notify_all(); @@ -188,6 +193,9 @@ void server_queue::yield_to_queue(std::function && work) { QUE_DBG("%s", "yielding to queue\n"); + // tasks declined while the work is running + std::deque unhandled; + { std::unique_lock lock(mutex_tasks); GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested"); @@ -200,11 +208,11 @@ void server_queue::yield_to_queue(std::function && work) { } while (true) { - // note: on terminate, this becomes a no-op and we simply keep waiting for the work to - // finish, we cannot return early because work() borrows the caller's stack - process_new_tasks(); + // note: on terminate this is a no-op, but we still wait for the work to finish + process_new_tasks(&unhandled); std::unique_lock lock(mutex_tasks); + // declined tasks are kept in unhandled, so a non-empty queue always has something new condition_tasks.wait(lock, [&]{ return !worker.busy || (running && !queue_tasks.empty()); }); @@ -214,14 +222,21 @@ void server_queue::yield_to_queue(std::function && work) { } { - // make sure to avoid idle timeout here std::unique_lock lock(mutex_tasks); + + // put the declined tasks back, keeping their order + while (!unhandled.empty()) { + queue_tasks.push_front(std::move(unhandled.back())); + unhandled.pop_back(); + } + + // make sure to avoid idle timeout here time_last_task = ggml_time_ms(); } QUE_DBG("%s", "done yielding to queue\n"); - // the worker thread is idle now, so we can safely access worker.exception + // the worker is idle now, safe to read worker.exception if (worker.exception) { std::exception_ptr exception = nullptr; std::swap(exception, worker.exception); diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index c36a8cee99..b207971a82 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -39,7 +39,7 @@ private: worker_t worker; // callback functions - std::function callback_new_task; + std::function callback_new_task; std::function callback_update_slots; std::function callback_sleeping_state; @@ -94,6 +94,8 @@ public: // 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 + // + // tasks declined by callback_new_task are put back in the queue once this returns void yield_to_queue(std::function && work); // for metrics @@ -107,7 +109,9 @@ public: // // Register function to process a new task - void on_new_task(std::function callback) { + // it returns false to decline the task, which is only allowed while yielding + // on decline, the task must be left untouched, it is put back in the queue later + void on_new_task(std::function callback) { callback_new_task = std::move(callback); } @@ -136,7 +140,8 @@ private: // process all pending tasks in the queue // returns true if the queue is terminated, false if there is no more task to process - bool process_new_tasks(); + // declined tasks are moved to unhandled, which must be set when called while yielding + bool process_new_tasks(std::deque * unhandled = nullptr); // for worker_t void worker_loop(); diff --git a/tools/server/tests/unit/test_metrics.py b/tools/server/tests/unit/test_metrics.py index 10cfc424b1..dfc8500c7a 100644 --- a/tools/server/tests/unit/test_metrics.py +++ b/tools/server/tests/unit/test_metrics.py @@ -225,3 +225,39 @@ def test_metrics_embedding_prompt_is_counted(): assert metrics["llamacpp:prompt_tokens_total"][1] > 0 assert metrics["llamacpp:n_decode_total"][1] > 0 assert metrics["llamacpp:tokens_predicted_total"][1] == 0 + + +def test_metrics_served_while_decoding(): + global server + server.server_slots = True + server.start() + + # the decode runs on a worker thread, so the queue can still answer metrics tasks + # ignore_eos keeps the slot busy for the whole request + stream = server.make_stream_request("POST", "/completion", data={ + "prompt": "the quick brown fox jumps over the lazy dog " * 4, + "n_predict": 200, + "ignore_eos": True, + "stream": True, + }) + next(stream) # the first token is out, the slot is decoding from now on + + # short timeouts: if the queue stops answering, fail fast instead of hanging for 10 minutes + res = server.make_request("GET", "/slots", timeout=30) + assert res.status_code == 200 + assert any(slot["is_processing"] for slot in res.body) + + res = server.make_request("GET", "/metrics", timeout=30) + assert res.status_code == 200 + assert isinstance(res.body, str) + assert parse_metrics(res.body)["llamacpp:requests_processing"][1] == 1 + + # a completion is not safe to run during a decode, it is declined and must not be lost + res = server.make_request("POST", "/completion", data={"prompt": "I believe", "n_predict": 4}, timeout=30) + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == 4 + + last = None + for chunk in stream: + last = chunk + assert last is not None and last["stop"] is True