Felipe97/llama-cpp-compiled

1#pragma once2 3#include "common.h"4#include "download.h"5#include "preset.h"6#include "server-common.h"7#include "server-http.h"8#include "server-queue.h"9 10#include <mutex>11#include <condition_variable>12#include <thread>13#include <functional>14#include <memory>15#include <optional>16#include <set>17#include <string>18#include <unordered_map>19 20/**21 * state diagram:22 *23 * DOWNLOADING ──► DOWNLOADED ──► (replaced by new instance)24 *25 * UNLOADED ──► LOADING ──► LOADED ◄──── SLEEPING26 *  ▲            │            │               ▲27 *  └───failed───┘            │               │28 *  ▲                         └──sleeping─────┘29 *  └────────unloaded─────────┘30 */31enum server_model_status {32    // TODO: also add downloading state when the logic is added33    SERVER_MODEL_STATUS_DOWNLOADING,34    SERVER_MODEL_STATUS_DOWNLOADED,35    SERVER_MODEL_STATUS_UNLOADED,36    SERVER_MODEL_STATUS_LOADING,37    SERVER_MODEL_STATUS_LOADED,38    SERVER_MODEL_STATUS_SLEEPING39};40 41enum server_model_source {42    SERVER_MODEL_SOURCE_PRESET,43    SERVER_MODEL_SOURCE_MODELS_DIR,44    SERVER_MODEL_SOURCE_CACHE,45};46 47enum server_child_mode {48    SERVER_CHILD_MODE_NORMAL,   // load the model and run normally49    SERVER_CHILD_MODE_DOWNLOAD, // download the model and exit50};51 52static std::string server_model_status_to_string(server_model_status status) {53    switch (status) {54        case SERVER_MODEL_STATUS_DOWNLOADING: return "downloading";55        case SERVER_MODEL_STATUS_DOWNLOADED:  return "downloaded";56        case SERVER_MODEL_STATUS_UNLOADED:    return "unloaded";57        case SERVER_MODEL_STATUS_LOADING:     return "loading";58        case SERVER_MODEL_STATUS_LOADED:      return "loaded";59        case SERVER_MODEL_STATUS_SLEEPING:    return "sleeping";60        default:                              return "unknown";61    }62}63 64static std::string server_model_source_to_string(server_model_source source) {65    switch (source) {66        case SERVER_MODEL_SOURCE_PRESET:     return "preset";67        case SERVER_MODEL_SOURCE_MODELS_DIR: return "models_dir";68        case SERVER_MODEL_SOURCE_CACHE:      return "cache";69        default:                             return "unknown";70    }71}72 73struct server_model_meta {74    server_model_source source = SERVER_MODEL_SOURCE_CACHE;75    common_preset preset;76    std::string name;77    std::set<std::string> aliases; // additional names that resolve to this model78    std::set<std::string> tags;    // informational tags, not used for routing79    int port = 0;80    server_model_status status = SERVER_MODEL_STATUS_UNLOADED;81    int64_t last_used = 0; // for LRU unloading82    std::vector<std::string> args; // args passed to the model instance, will be populated by render_args()83    json loaded_info; // info to be reflected via /v1/models endpoint ; if in DOWNLOADING state, it should contain download progress info84    json progress; // reflect load or download progress info, if any85    int exit_code = 0; // exit code of the model instance process (only valid if status == FAILED)86    int stop_timeout = 0; // seconds to wait before force-killing the model instance during shutdown87    mtmd_caps multimodal; // multimodal capabilities88    bool hidden = false; // hidden from GET /models, but still accept if requested89 90    bool is_ready() const {91        return status == SERVER_MODEL_STATUS_LOADED;92    }93 94    bool is_running() const {95        return status == SERVER_MODEL_STATUS_LOADED || status == SERVER_MODEL_STATUS_LOADING || status == SERVER_MODEL_STATUS_SLEEPING;96    }97 98    bool is_ready_or_sleep() const {99        return status == SERVER_MODEL_STATUS_LOADED || status == SERVER_MODEL_STATUS_SLEEPING;100    }101 102    bool is_failed() const {103        return status == SERVER_MODEL_STATUS_UNLOADED && exit_code != 0;104    }105 106    void update_args(common_preset_context & ctx_presets, std::string bin_path);107    void update_caps();108};109 110struct server_models_routes;111struct server_lru_sched; // defined in server-models.cpp112struct server_monitor;   // defined in server-models.cpp113 114struct server_models {115    friend struct server_models_routes;116    friend struct server_lru_sched;117    friend struct server_monitor;118 119private:120    struct instance_t {121        std::shared_ptr<server_subproc> subproc; // shared with the monitor thread122        server_model_meta meta;123        int req_count = 0; // number of active proxy requests124 125        // ask the child to exit (it handles the command on its stdin, see server_child::setup)126        void request_exit() const;127    };128 129    std::mutex mutex;130    std::condition_variable cv;131    std::map<std::string, instance_t> mapping;132 133    // models asked to stop, still counted as running until the monitor records their exit134    std::set<std::string> stopping_models;135 136    // set to true while load_models() is executing a reload; load() will wait until clear137    bool is_reloading = false;138 139    // if true, the next get_meta() will trigger a reload of model list140    bool need_reload = false;141 142    // models marked with load-on-startup, unset once load_startup_models() drains it143    // no value means the startup phase is over, so a reload must not queue anything144    std::optional<std::vector<std::string>> startup_models{std::in_place};145 146    // conv_id -> model name that currently serves its stream session, lets the resumable stream147    // routes go straight to the owning child instead of polling every one. populated when148    // proxy_request forwards a POST carrying an X-Conversation-Id. best effort: a stale entry just149    // makes the child answer not found and the client recovers. owns its lock, one mutex per struct150    struct conv_model_tracker {151        // returns the ticket of this registration, 0 when nothing was registered. erasing or152        // replacing the entry invalidates the ticket, which is how a stop cancels a request153        // parked in the model load wait154        uint64_t remember(const std::string & conv_id, const std::string & model) {155            if (conv_id.empty() || model.empty()) {156                return 0;157            }158            std::lock_guard<std::mutex> lock(mu);159            uint64_t ticket = next_ticket++;160            map[conv_id] = { model, ticket };161            return ticket;162        }163 164        // false means a stop erased the entry or a newer request replaced it165        bool alive(const std::string & conv_id, uint64_t ticket) {166            std::lock_guard<std::mutex> lock(mu);167            auto it = map.find(conv_id);168            return it != map.end() && it->second.ticket == ticket;169        }170 171        std::optional<std::string> lookup(const std::string & conv_id) {172            if (conv_id.empty()) {173                return std::nullopt;174            }175            std::lock_guard<std::mutex> lock(mu);176            auto it = map.find(conv_id);177            if (it == map.end()) {178                return std::nullopt;179            }180            return it->second.model;181        }182 183        void forget(const std::string & conv_id) {184            if (conv_id.empty()) {185                return;186            }187            std::lock_guard<std::mutex> lock(mu);188            map.erase(conv_id);189        }190 191      private:192        struct entry_t {193            std::string model;194            uint64_t    ticket;195        };196        std::mutex                               mu;197        uint64_t                                 next_ticket = 1;198        std::unordered_map<std::string, entry_t> map;199    };200 201    common_preset_context ctx_preset;202 203    common_params base_params;204    std::string bin_path;205    std::vector<std::string> base_env;206    common_preset base_preset; // base preset from llama-server CLI args207 208    // queue of requests waiting for a models_max slot209    std::unique_ptr<server_lru_sched> sched;210 211    // if true, add some delay to simulate works (useful for testing)212    bool debug_fake_timing = false;213 214    void update_meta(const std::string & name, const server_model_meta & meta);215 216    // unload least recently used models if the limit is reached217    void unload_lru();218 219    // not thread-safe, caller must hold mutex220    void add_model(server_model_meta && meta);221 222    // ask the monitor to stop a running instance; send_exit is false for a child that was already force-killed223    // not thread-safe, caller must hold mutex224    void request_stop(const std::string & name, bool send_exit = true);225 226    // called by the monitor once a child exited and was reaped227    void on_child_exit(const std::string & name, const std::shared_ptr<server_subproc> & proc, server_child_mode mode, int exit_code);228 229    // notify SSE clients230    void notify_sse(const std::string & event, const std::string & model_id, const json & data = nullptr);231 232public:233    // conv_id -> model tracker for the resumable stream routes, owns its lock234    conv_model_tracker conv_models;235 236    server_models(const common_params & params, int argc, char ** argv);237    ~server_models();238 239    server_response sse; // for real-time updates via SSE endpoint240 241    // (re-)load the list of models from various sources and prepare the metadata mapping242    // - if this is called the first time, simply populate the metadata243    // - if this is called subsequently (e.g. when refreshing from disk):244    //   - if a model is running but updated or removed from the source, it will be unloaded245    //   - if a model is not running, it will be added or updated according to the source246    void load_models();247 248    // lazy-load startup_models, to be called after main() setup phase249    void load_startup_models();250 251    // check if a model instance exists (thread-safe)252    bool has_model(const std::string & name);253 254    // return a copy of model metadata (thread-safe)255    std::optional<server_model_meta> get_meta(const std::string & name);256 257    // return a copy of all model metadata (thread-safe)258    std::vector<server_model_meta> get_all_meta();259 260    struct load_options {261        server_child_mode mode = SERVER_CHILD_MODE_NORMAL;262        // used for spawning a downloading child process263        std::optional<server_model_meta> custom_meta = std::nullopt;264    };265 266    // load and unload model instances267    // these functions are thread-safe268    void load(const std::string & name);269    void load(const std::string & name, const load_options & opts);270    void unload(const std::string & name);271    void unload_all();272 273    struct update_status_args {274        server_model_status status;275        int exit_code = 0; // only valid if status == UNLOADED276        json loaded_info = nullptr;277        json progress = nullptr;278    };279    // update the status of a model instance (thread-safe)280    // also send SSE notification to /models/sse endpoint281    void update_status(const std::string & name, const update_status_args & args);282    void update_download_progress(const std::string & name, const common_download_progress & progress, bool done, bool ok = true);283 284    // remove a cache model from disk and update the list (thread-safe)285    // note: only cache models can be removed; returns false if the model doesn't exist or is not a cache model286    bool remove(const std::string & name);287 288    // wait until the model instance is fully loaded (thread-safe)289    // note: predicate is called while holding the lock290    // return when the model no longer in "loading" state291    void wait(const std::string & name, std::function<bool(const server_model_meta &)> predicate);292    void wait(std::unique_lock<std::mutex> & lk, const std::string & name, std::function<bool(const server_model_meta &)> predicate);293 294    // ensure the model is in ready state (thread-safe)295    // return false if model is ready296    // otherwise, load the model and blocking wait until it's ready, then return true (meta may need to be refreshed)297    // if models_max is reached, the request waits in a queue until a slot frees up298    // throws if the load fails, or if should_stop fires while waiting299    bool ensure_model_ready(const std::string & name, const std::function<bool()> & should_stop = nullptr);300 301    // proxy an HTTP request to the model instance302    server_http_res_ptr proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used, bool detached = false);303 304    // handle message sent from server_child::notify_to_router()305    // raw input must starts with CMD_CHILD_TO_ROUTER_STATE, followed by a JSON string306    // called from the monitor thread307    // payload per state:308    //     state = loading     -> payload = {} (TODO: add progress info)309    //     state = ready       -> payload = model_info (json), or {} if wakeup from sleeping310    //     state = sleeping    -> payload = {}311    void handle_child_state(const std::string & name, const std::string & raw_input);312 313private:314    // one thread watching every child; keep last, the destructor joins the thread315    std::unique_ptr<server_monitor> monitor;316};317 318struct server_child {319    // serializes the notify_to_router writes320    std::mutex mtx_stdout;321    std::atomic<bool> is_finished_downloading = false; // set by run_download322 323    // return true if the current process is a child server instance324    bool is_child();325    server_child_mode get_mode();326    int run_download(common_params & params);327 328    // register the shutdown_handler to be called by the router329    // return the monitoring thread (to be joined by the caller)330    std::thread setup(const std::function<void(int)> & shutdown_handler);331 332    // notify router server for status changes (e.g. loading, downloading, sleeping, etc.)333    // message will be handled by server_models::handle_child_state() on the router side334    void notify_to_router(const std::string & state_name, const json & payload);335};336 337struct server_models_routes {338    common_params params;339    json ui_settings = json::object();     // Primary: new name340    std::atomic<bool> stopping = false;    // for graceful disconnecting SSE clients during shutdown341    server_models models;342    server_models_routes(const common_params & params, int argc, char ** argv)343            : params(params), models(params, argc, argv) {344        const std::string & cfg = this->params.ui_config_json;345        if (!cfg.empty()) {346            try {347                json json_settings = json::parse(cfg);348                ui_settings = json_settings;349            } catch (const std::exception & e) {350                LOG_ERR("%s: failed to parse UI config: %s\n", __func__, e.what());351                throw;352            }353        }354        init_routes();355    }356 357    void init_routes();358    // handlers using lambda function, so that they can capture `this` without `std::bind`359    server_http_context::handler_t get_router_props;360    server_http_context::handler_t proxy_get;361    server_http_context::handler_t proxy_post;362    server_http_context::handler_t get_router_models;363    server_http_context::handler_t post_router_models_load;364    server_http_context::handler_t post_router_models_unload;365    // management API366    server_http_context::handler_t get_router_models_sse;367    server_http_context::handler_t post_router_models;368    server_http_context::handler_t del_router_models;369 370    // router side handlers for the resumable streaming routes. each resolves the child that owns371    // a conversation through the conv_id -> model map, no probing or fan out372    server_http_context::handler_t router_stream_get;373    server_http_context::handler_t router_streams_lookup;374    server_http_context::handler_t router_stream_delete;375};376 377/**378 * A simple HTTP proxy that forwards requests to another server379 * and relays the responses back.380 */381struct server_http_proxy : server_http_res {382    std::function<void()> cleanup = nullptr;383    server_http_proxy(const std::string & method,384                      const std::string & scheme,385                      const std::string & host,386                      int port,387                      const std::string & path,388                      const std::map<std::string, std::string> & headers,389                      const std::string & body,390                      const std::map<std::string, uploaded_file> & files,391                      const std::function<bool()> should_stop,392                      int32_t timeout_read,393                      int32_t timeout_write394                      );395    ~server_http_proxy() {396        if (cleanup_pipes) {397            cleanup_pipes();398        }399        if (cleanup) {400            cleanup();401        }402    }403private:404    std::function<void()> cleanup_pipes = nullptr;405    std::thread thread;406    struct msg_t {407        std::map<std::string, std::string> headers;408        int status = 0;409        std::string data;410        std::string content_type;411    };412};413