echodict/llama.cpp

1#pragma once2 3#include "server-task.h"4 5#include <condition_variable>6#include <deque>7#include <mutex>8#include <vector>9#include <unordered_set>10 11// struct for managing server tasks12// in most cases, use server_response_reader to post new tasks and retrieve results13struct server_queue {14private:15    int id = 0;16    bool running  = false;17    bool sleeping = false;18    bool req_stop_sleeping = false;19    int64_t time_last_task = 0;20 21    // queues22    std::deque<server_task> queue_tasks;23    std::deque<server_task> queue_tasks_deferred;24 25    std::mutex mutex_tasks;26    std::condition_variable condition_tasks;27 28    // callback functions29    std::function<void(server_task &&)> callback_new_task;30    std::function<void(void)>           callback_update_slots;31    std::function<void(bool)>           callback_sleeping_state;32 33public:34    // Add a new task to the end of the queue35    int post(server_task && task, bool front = false);36 37    // multi-task version of post()38    int post(std::vector<server_task> && tasks, bool front = false);39 40    // Add a new task, but defer until one slot is available41    void defer(server_task && task);42 43    // Get the next id for creating a new task44    int get_new_id();45 46    // Call when the state of one slot is changed, it will move one task from deferred to main queue47    // prioritize tasks that use the specified slot (otherwise, pop the first deferred task)48    void pop_deferred_task(int id_slot);49 50    // if sleeping, request exiting sleep state and wait until it is done51    // returns immediately if not sleeping52    void wait_until_no_sleep();53 54    bool is_sleeping() {55        std::unique_lock<std::mutex> lock(mutex_tasks);56        return sleeping;57    }58 59    // end the start_loop routine60    void terminate();61 62    /**63     * Main loop consists of these steps:64     * - Wait until a new task arrives65     * - Process the task (i.e. maybe copy data into slot)66     * - Check if multitask is finished67     * - Update all slots68     *69     * Sleeping procedure (disabled if idle_sleep_ms < 0):70     * - If there is no task after idle_sleep_ms, enter sleeping state71     * - Call callback_sleeping_state(true)72     * - Wait until req_stop_sleeping is set to true73     * - Call callback_sleeping_state(false)74     * - Exit sleeping state75     */76    void start_loop(int64_t idle_sleep_ms = -1);77 78    // for metrics79    size_t queue_tasks_deferred_size() {80        std::unique_lock<std::mutex> lock(mutex_tasks);81        return queue_tasks_deferred.size();82    }83 84    //85    // Functions below are not thread-safe, must only be used before start_loop() is called86    //87 88    // Register function to process a new task89    void on_new_task(std::function<void(server_task &&)> callback) {90        callback_new_task = std::move(callback);91    }92 93    // Register the function to be called when all slots data is ready to be processed94    void on_update_slots(std::function<void(void)> callback) {95        callback_update_slots = std::move(callback);96    }97 98    // Register callback for sleeping state change; multiple callbacks are allowed99    // note: when entering sleeping state, the callback is called AFTER sleeping is set to true100    //       when leaving sleeping state, the callback is called BEFORE sleeping is set to false101    void on_sleeping_state(std::function<void(bool)> callback) {102        if (callback_sleeping_state) {103            auto prev_callback = std::move(callback_sleeping_state);104            callback_sleeping_state = [prev_callback, callback](bool sleeping) {105                prev_callback(sleeping);106                callback(sleeping);107            };108        } else {109            callback_sleeping_state = std::move(callback);110        }111    }112 113private:114    void cleanup_pending_task(int id_target);115};116 117// struct for managing server responses118// in most cases, use server_response_reader to retrieve results119struct server_response {120private:121    bool running = true;122 123    // for keeping track of all tasks waiting for the result124    std::unordered_set<int> waiting_task_ids;125 126    // the main result queue (using ptr for polymorphism)127    std::vector<server_task_result_ptr> queue_results;128 129    std::mutex mutex_results;130    std::condition_variable condition_results;131 132public:133    // add the id_task to the list of tasks waiting for response134    void add_waiting_task_id(int id_task);135 136    void add_waiting_task_ids(const std::unordered_set<int> & id_tasks);137 138    // when the request is finished, we can remove task associated with it139    void remove_waiting_task_id(int id_task);140 141    // remove multiple tasks from waiting list142    void remove_waiting_task_ids(const std::unordered_set<int> & id_tasks);143 144    // This function blocks the thread until there is a response for one of the id_tasks145    server_task_result_ptr recv(const std::unordered_set<int> & id_tasks);146 147    // same as recv(), but have timeout in seconds148    // if timeout is reached, nullptr is returned149    server_task_result_ptr recv_with_timeout(const std::unordered_set<int> & id_tasks, int timeout);150 151    // single-task version of recv()152    server_task_result_ptr recv(int id_task);153 154    // Send a new result to a waiting id_task155    void send(server_task_result_ptr && result);156 157    // terminate the waiting loop158    void terminate();159};160 161// utility class to make working with server_queue and server_response easier162// it provides a generator-like API for server responses163// support pooling connection state and aggregating multiple results164struct server_response_reader {165    std::unordered_set<int> id_tasks;166    server_queue & queue_tasks;167    server_response & queue_results;168    size_t received_count = 0;169    bool cancelled = false;170    int polling_interval_seconds;171 172    // tracking generation state and partial tool calls173    // only used by streaming completions174    std::vector<task_result_state> states;175 176    // should_stop function will be called each polling_interval_seconds177    server_response_reader(server_queue & queue_tasks, server_response & queue_results, int polling_interval_seconds)178        : queue_tasks(queue_tasks), queue_results(queue_results), polling_interval_seconds(polling_interval_seconds) {}179    ~server_response_reader() {180        stop();181    }182 183    int get_new_id() {184        return queue_tasks.get_new_id();185    }186 187    // if front = true, the task will be posted to the front of the queue (high priority)188    void post_task(server_task && task, bool front = false);189    void post_tasks(std::vector<server_task> && tasks, bool front = false);190    bool has_next() const;191 192    // return nullptr if should_stop() is true before receiving a result193    // note: if one error is received, it will stop further processing and return error result194    server_task_result_ptr next(const std::function<bool()> & should_stop);195 196    struct batch_response {197        bool is_terminated = false; // if true, indicates that processing was stopped before all results were received198        std::vector<server_task_result_ptr> results;199        server_task_result_ptr error; // nullptr if no error200    };201    // aggregate multiple results202    batch_response wait_for_all(const std::function<bool()> & should_stop);203 204    void stop();205};206