server_queue::worker

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,121 @@ void server_queue::terminate() {
     condition_tasks.notify_all();
 }
 
+bool server_queue::process_new_tasks() {
+    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);
+        callback_new_task(std::move(task));
+    }
+}
+
+void server_queue::worker_loop() {
+    while (true) {
+        std::function<void()> work;
+        {
+            std::unique_lock<std::mutex> lock(worker.mutex);
+            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
+        try {
+            work();
+        } catch (...) {
+            worker.exception = std::current_exception();
+        }
+
+        // signal completion to the thread waiting in yield_to_queue()
+        std::unique_lock<std::mutex> lock(mutex_tasks);
+        worker.busy = false;
+        condition_tasks.notify_all();
+    }
+}
+
+void server_queue::worker_stop() {
+    if (!worker.thread.joinable()) {
+        return;
+    }
+    {
+        std::unique_lock<std::mutex> lock(worker.mutex);
+        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() must be called from the start_loop() thread");
+
+    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;
+    }
+    {
+        std::unique_lock<std::mutex> lock(worker.mutex);
+        worker.work = std::move(work);
+        worker.cv.notify_one();
+    }
+
+    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();
+
+        std::unique_lock<std::mutex> lock(mutex_tasks);
+        condition_tasks.wait(lock, [&]{
+            return !worker.busy || (running && !queue_tasks.empty());
+        });
+        if (!worker.busy) {
+            break;
+        }
+    }
+
+    {
+        // make sure to avoid idle timeout here
+        std::unique_lock<std::mutex> lock(mutex_tasks);
+        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
+    if (worker.exception) {
+        std::exception_ptr exception = nullptr;
+        std::swap(exception, worker.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 +250,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()) {
+            break; // terminate
         }
+
         // all tasks in the current loop is processed, slots data is now ready
         QUE_DBG("%s", "update slots\n");
 
@@ -206,6 +304,8 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
             }
         }
     }
+
+    worker_stop();
 }
 
 void server_queue::cleanup_pending_task(int id_target) {

tools/server/server-queue.h

@@ -5,6 +5,7 @@
 #include <condition_variable>
 #include <deque>
 #include <mutex>
+#include <thread>
 #include <vector>
 #include <unordered_set>
 
@@ -25,12 +26,26 @@ private:
     std::mutex mutex_tasks;
     std::condition_variable condition_tasks;
 
+    // used by yield_to_queue
+    struct worker_t {
+        std::thread thread;
+        std::mutex  mutex;
+        std::condition_variable cv;
+        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; // guarded by mutex_tasks (not by worker_t::mutex)
+    };
+    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;
 
 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 +90,12 @@ 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
+    void yield_to_queue(std::function<void()> && work);
+
     // for metrics
     size_t queue_tasks_deferred_size() {
         std::unique_lock<std::mutex> lock(mutex_tasks);
@@ -112,6 +133,14 @@ 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
+    bool process_new_tasks();
+
+    // for worker_t
+    void worker_loop();
+    void worker_stop();
 };
 
 // struct for managing server responses