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-queue.cpp

@@ -4,6 +4,7 @@
 #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__)
@@ -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<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
@@ -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<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");
 
@@ -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());
 }
 
 //