call llama_decode inside yield_to_queue

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<server_prompt_cache> 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<server_slot> & slots, std::function<void(server_slot &)> 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) {
             {

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<server_task> * unhandled) {
     while (true) {
         std::unique_lock<std::mutex> 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<std::mutex> lock(mutex_tasks);
         worker.busy = false;
         condition_tasks.notify_all();
@@ -188,6 +193,9 @@ void server_queue::yield_to_queue(std::function<void()> && work) {
 
     QUE_DBG("%s", "yielding to queue\n");
 
+    // tasks declined while the work is running
+    std::deque<server_task> unhandled;
+
     {
         std::unique_lock<std::mutex> 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<void()> && 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<std::mutex> 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<void()> && work) {
     }
 
     {
-        // make sure to avoid idle timeout here
         std::unique_lock<std::mutex> 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);

tools/server/server-queue.h

@@ -39,7 +39,7 @@ private:
     worker_t worker;
 
     // callback functions
-    std::function<void(server_task &&)> callback_new_task;
+    std::function<bool(server_task &&)> callback_new_task;
     std::function<void(void)>           callback_update_slots;
     std::function<void(bool)>           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<void()> && work);
 
     // for metrics
@@ -107,7 +109,9 @@ public:
     //
 
     // Register function to process a new task
-    void on_new_task(std::function<void(server_task &&)> 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<bool(server_task &&)> 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<server_task> * unhandled = nullptr);
 
     // for worker_t
     void worker_loop();

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