Felipe97/llama-cpp-compiled

1#pragma once2 3#include "server-task.h"4 5#include <condition_variable>6#include <deque>7#include <exception>8#include <mutex>9#include <thread>10#include <vector>11#include <unordered_set>12 13// struct for managing server tasks14// in most cases, use server_response_reader to post new tasks and retrieve results15struct server_queue {16private:17    int id = 0;18    bool running  = false;19    bool sleeping = false;20    bool req_stop_sleeping = false;21    int64_t time_last_task = 0;22 23    // queues24    std::deque<server_task> queue_tasks;25    std::deque<server_task> queue_tasks_deferred;26    // tasks declined while yielding, put back in queue_tasks once the yield is done27    // note: kept as a member so that cleanup_pending_task() can also reach them28    std::deque<server_task> queue_tasks_unhandled;29 30    std::mutex mutex_tasks;31    std::condition_variable condition_tasks;32 33    // used by yield_to_queue, all fields are guarded by mutex_tasks34    struct worker_t {35        std::thread             thread;36        std::condition_variable cv;        // the worker sleeps on this until a yield starts37        std::exception_ptr      exception; // exception thrown while processing tasks, if any38        bool stop     = false;39        bool busy     = false; // set by yield_to_queue(), cleared by the worker once it is done processing tasks40        bool yielding = false; // work() is still running on the start_loop() thread41    };42    worker_t worker;43 44    // callback functions45    std::function<bool(server_task &&, bool)> callback_new_task;46    std::function<void(void)>                 callback_update_slots;47    std::vector<std::function<void(bool)>>    callback_sleeping_state;48 49public:50    ~server_queue() { worker_stop(); }51 52    // Add a new task to the end of the queue53    int post(server_task && task, bool front = false);54 55    // multi-task version of post()56    int post(std::vector<server_task> && tasks, bool front = false);57 58    // Add a new task, but defer until one slot is available59    void defer(server_task && task);60 61    // Get the next id for creating a new task62    int get_new_id();63 64    // Call when the state of one slot is changed, it will move one task from deferred to main queue65    // prioritize tasks that use the specified slot (otherwise, pop the first deferred task)66    void pop_deferred_task(int id_slot);67 68    // if sleeping, request exiting sleep state and wait until it is done69    // returns immediately if not sleeping70    void wait_until_no_sleep();71 72    bool is_sleeping() {73        std::unique_lock<std::mutex> lock(mutex_tasks);74        return sleeping;75    }76 77    // end the start_loop routine78    void terminate();79 80    /**81     * Main loop consists of these steps:82     * - Wait until a new task arrives83     * - Process the task (i.e. maybe copy data into slot)84     * - Check if multitask is finished85     * - Update all slots86     *87     * Sleeping procedure (disabled if idle_sleep_ms < 0):88     * - If there is no task after idle_sleep_ms, enter sleeping state89     *   note: metrics tasks are processed as usual, but do not reset the idle timer90     * - Call callback_sleeping_state(true)91     * - Wait until req_stop_sleeping is set to true92     * - Call callback_sleeping_state(false)93     * - Exit sleeping state94     */95    void start_loop(int64_t idle_sleep_ms = -1);96 97    // while waiting for work() to finish, run process_new_tasks on the worker thread98    // returns once work() is done (may throw exceptions)99    // must be called from start_loop() thread (ideally inside callback_update_slots)100    // use case: return metrics while encode/decode is running101    // ref: https://github.com/ggml-org/llama.cpp/pull/27041102    //103    // tasks declined by callback_new_task are put back in the queue once this returns104    void yield_to_queue(std::function<void()> && work);105 106    // for metrics107    size_t queue_tasks_deferred_size() {108        std::unique_lock<std::mutex> lock(mutex_tasks);109        return queue_tasks_deferred.size();110    }111 112    //113    // Functions below are not thread-safe, must only be used before start_loop() is called114    //115 116    // Register function to process a new task117    // the second argument tells whether the queue is currently yielding (see yield_to_queue)118    // only then may the callback return false to decline the task, and it must leave it119    // untouched, so that it can be put back in the queue later120    // note: while yielding, the callback runs on worker thread, not main thread121    void on_new_task(std::function<bool(server_task &&, bool)> callback) {122        callback_new_task = std::move(callback);123    }124 125    // Register the function to be called when all slots data is ready to be processed126    void on_update_slots(std::function<void(void)> callback) {127        callback_update_slots = std::move(callback);128    }129 130    // Register callback for sleeping state change; multiple callbacks are allowed131    // for example: register order cb0, cb1, cb2132    // entering sleep: queue.sleeping = true --> cb0(true) --> cb1(true) --> cb2(true)133    // leaving sleep: cb2(false) --> cb1(false) --> cb0(false) --> queue.sleeping = false134    // note: caller will hold mutex_tasks while calling the callbacks135    void on_sleeping_state(std::function<void(bool)> callback) {136        callback_sleeping_state.push_back(std::move(callback));137    }138 139private:140    void cleanup_pending_task(int id_target);141 142    // process all pending tasks in the queue143    // returns true if the queue is terminated, false if there is no more task to process144    // while yielding, declined tasks are moved to queue_tasks_unhandled145    bool process_new_tasks(bool is_yielding);146 147    // for worker_t148    void worker_loop();149    void worker_stop();150};151 152// struct for managing server responses153// in most cases, use server_response_reader to retrieve results154struct server_response {155private:156    bool running = true;157 158    // for keeping track of all tasks waiting for the result159    std::unordered_set<int> waiting_task_ids;160 161    // the main result queue (using ptr for polymorphism)162    std::vector<server_task_result_ptr> queue_results;163 164    std::mutex mutex_results;165    std::condition_variable condition_results;166 167public:168    // add the id_task to the list of tasks waiting for response169    void add_waiting_task_id(int id_task);170 171    void add_waiting_task_ids(const std::unordered_set<int> & id_tasks);172 173    // when the request is finished, we can remove task associated with it174    void remove_waiting_task_id(int id_task);175 176    // remove multiple tasks from waiting list177    void remove_waiting_task_ids(const std::unordered_set<int> & id_tasks);178 179    // This function blocks the thread until there is a response for one of the id_tasks180    server_task_result_ptr recv(const std::unordered_set<int> & id_tasks);181 182    // same as recv(), but have timeout in seconds183    // if timeout is reached, nullptr is returned184    server_task_result_ptr recv_with_timeout(const std::unordered_set<int> & id_tasks, int timeout);185 186    // single-task version of recv()187    server_task_result_ptr recv(int id_task);188 189    // Send a new result to a waiting id_task190    void send(server_task_result_ptr && result);191 192    // broadcast a new result to all waiting tasks193    // (used by router mode)194    void broadcast(server_task_result_ptr && result);195 196    // terminate the waiting loop197    void terminate();198};199 200// RAII wrapper to make working with server_queue and server_response easier201// it provides a generator-like API for server responses202// support pooling connection state and aggregating multiple results203struct server_response_reader {204    std::unordered_set<int> id_tasks;205    server_queue & queue_tasks;206    server_response & queue_results;207    size_t received_count = 0;208    bool cancelled = false;209    int polling_interval_seconds;210 211    // tracking generation state and partial tool calls212    // only used by streaming completions213    std::vector<task_result_state> states;214 215    // should_stop function will be called each polling_interval_seconds216    server_response_reader(server_queue & queue_tasks, server_response & queue_results, int polling_interval_seconds)217        : queue_tasks(queue_tasks), queue_results(queue_results), polling_interval_seconds(polling_interval_seconds) {}218    ~server_response_reader() {219        stop();220    }221 222    int get_new_id() {223        return queue_tasks.get_new_id();224    }225 226    // if front = true, the task will be posted to the front of the queue (high priority)227    void post_task(server_task && task, bool front = false);228    void post_tasks(std::vector<server_task> && tasks, bool front = false);229    bool has_next() const;230 231    // return nullptr if should_stop() is true before receiving a result232    // note: if one error is received, it will stop further processing and return error result233    server_task_result_ptr next(const std::function<bool()> & should_stop);234 235    struct batch_response {236        bool is_terminated = false; // if true, indicates that processing was stopped before all results were received237        std::vector<server_task_result_ptr> results;238        server_task_result_ptr error; // nullptr if no error239    };240    // aggregate multiple results241    batch_response wait_for_all(const std::function<bool()> & should_stop);242 243    void stop();244};245