char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Bartowskiserver: add ctx-per-slot (--kv-unified-per-slot) (#24124)18443257a

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
23.0 KiB562 linesraw
1#include "server-context.h"
2#include "server-http.h"
3#include "server-models.h"
4#include "server-cors-proxy.h"
5#include "server-stream.h"
6#include "server-tools.h"
7
8#include "arg.h"
9#include "build-info.h"
10#include "common.h"
11#include "fit.h"
12#include "llama.h"
13#include "log.h"
14
15#include <atomic>
16#include <clocale>
17#include <exception>
18#include <signal.h>
19#include <thread> // for std::thread::hardware_concurrency
20
21#if defined(_WIN32)
22#include <windows.h>
23#endif
24
25static std::function<void(int)> shutdown_handler;
26static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT;
27
28static inline void signal_handler(int signal) {
29    if (is_terminating.test_and_set()) {
30        // in case it hangs, we can force terminate the server by hitting Ctrl+C twice
31        // this is for better developer experience, we can remove when the server is stable enough
32        fprintf(stderr, "Received second interrupt, terminating immediately.\n");
33        exit(1);
34    }
35
36    shutdown_handler(signal);
37}
38
39// satisfies -Wmissing-declarations (used by llama command)
40int llama_server(int argc, char ** argv);
41
42// to be used via CLI (argc / argv are used by router mode only)
43int llama_server(common_params & params, int argc, char ** argv);
44void llama_server_terminate();
45void llama_server_terminate() {
46    if (shutdown_handler) {
47        shutdown_handler(0);
48    }
49}
50
51
52// wrapper function that handles exceptions and logs errors
53// this is to make sure handler_t never throws exceptions; instead, it returns an error response
54static server_http_context::handler_t ex_wrapper(server_http_context::handler_t func) {
55    return [func = std::move(func)](const server_http_req & req) -> server_http_res_ptr {
56        std::string message;
57        error_type error;
58        try {
59            return func(req);
60        } catch (const std::invalid_argument & e) {
61            // treat invalid_argument as invalid request (400)
62            error = ERROR_TYPE_INVALID_REQUEST;
63            message = e.what();
64        } catch (const std::exception & e) {
65            // treat other exceptions as server error (500)
66            error = ERROR_TYPE_SERVER;
67            message = e.what();
68        } catch (...) {
69            error = ERROR_TYPE_SERVER;
70            message = "unknown error";
71        }
72
73        auto res = std::make_unique<server_http_res>();
74        res->status = 500;
75        try {
76            json error_data = format_error_response(message, error);
77            res->status = json_value(error_data, "code", 500);
78            res->data = safe_json_to_str({{ "error", error_data }});
79            SRV_WRN("got exception: %s\n", res->data.c_str());
80        } catch (const std::exception & e) {
81            SRV_ERR("got another exception: %s | while handling exception: %s\n", e.what(), message.c_str());
82            res->data = "Internal Server Error";
83        }
84        return res;
85    };
86}
87
88int llama_server(int argc, char ** argv) {
89    std::setlocale(LC_NUMERIC, "C");
90
91#ifndef _WIN32
92    // Ignore SIGPIPE so the server does not crash if a child (MCP server, tools runtime) exits while we are writing to its stdin
93    signal(SIGPIPE, SIG_IGN);
94#endif
95
96    // own arguments required by this example
97    common_params params;
98
99    common_init();
100
101    // start the stream session manager GC right after common init, before any HTTP route can
102    // touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free
103    server_stream_session_manager_start();
104
105    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) {
106        return 1;
107    }
108
109    llama_backend_init();
110    llama_numa_init(params.numa);
111
112    return llama_server(params, argc, argv);
113}
114
115int llama_server(common_params & params, int argc, char ** argv) {
116    bool is_run_by_cli = (argv == nullptr);
117
118    common_models_handler models_handler;
119
120    // note: router mode also accepts -hf remote-preset, so we need to check that first
121    if (!is_run_by_cli && !params.model.hf_repo.empty()) {
122        try {
123            models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
124            if (common_models_handler_is_preset_repo(models_handler)) {
125                // apply the preset and start the server in router mode
126                common_models_handler_apply(models_handler, params);
127            }
128        } catch (const std::exception & e) {
129            SRV_ERR("failed to fetch model metadata: %s\n", e.what());
130            return 1;
131        }
132    }
133
134    // router server never loads a model and must not touch the GPU
135    const bool is_router_server = params.model.path.empty()
136                               && params.model.hf_repo.empty()
137                               && params.model.docker_repo.empty();
138
139    // skip device enumeration so the CUDA primary context stays uncreated
140    common_params_print_info(params, !is_router_server);
141
142    if (!is_router_server) {
143        // validate batch size for embeddings
144        // embeddings require all tokens to be processed in a single ubatch
145        // see https://github.com/ggml-org/llama.cpp/issues/12836
146        if (params.embedding && params.n_batch > params.n_ubatch) {
147            SRV_WRN("embeddings enabled with n_batch (%d) > n_ubatch (%d)\n", params.n_batch, params.n_ubatch);
148            SRV_WRN("setting n_batch = n_ubatch = %d to avoid assertion failure\n", params.n_ubatch);
149            params.n_batch = params.n_ubatch;
150        }
151
152        if (params.n_parallel < 0) {
153            SRV_TRC("%s", "n_parallel is set to auto, using n_parallel = 4 and kv_unified = true\n");
154
155            params.n_parallel = 4;
156            params.kv_unified = true;
157        }
158    }
159
160    // size the KV pool from --kv-unified-per-slot, unless the user pinned it with -c
161    // or with -c 0 for max context
162    const bool ctx_pool_auto_sized = params.kv_unified_per_slot > 0 &&
163                                     params.n_ctx == 0 &&
164                                     (uint32_t) params.fit_params_min_ctx != UINT32_MAX;
165
166    if (ctx_pool_auto_sized) {
167        params.n_ctx = params.n_parallel * params.kv_unified_per_slot;
168        SRV_INF("--kv-unified-per-slot: sizing KV pool to n_parallel * kv_unified_per_slot = %d * %d = %d\n", params.n_parallel,
169                params.kv_unified_per_slot, params.n_ctx);
170    }
171
172    // for consistency between server router mode and single-model mode, we set the same model name as alias
173    auto model_name = params.model.get_name();
174    if (params.model_alias.empty() && !model_name.empty()) {
175        params.model_alias.insert(model_name);
176    }
177
178    // note: this is guaranteed to out-live ctx_http and tools
179    server_mcp mcp_mgr;
180
181    // struct that contains llama context and inference
182    server_context ctx_server;
183
184    server_http_context ctx_http;
185    if (!ctx_http.init(params)) {
186        SRV_ERR("%s", "failed to initialize HTTP server\n");
187        return 1;
188    }
189
190    //
191    // Router
192    //
193
194    // register API routes
195    server_child child; // only used in non-router mode
196    server_routes routes(params, ctx_server);
197    server_tools tools;
198
199    std::optional<server_models_routes> models_routes{};
200    if (is_router_server) {
201        // setup server instances manager
202        try {
203            models_routes.emplace(params, argc, argv);
204        } catch (const std::exception & e) {
205            SRV_ERR("failed to initialize router models: %s\n", e.what());
206            return 1;
207        }
208
209        // proxy handlers
210        // note: routes.get_health stays the same
211        routes.get_metrics                 = models_routes->proxy_get;
212        routes.post_props                  = models_routes->proxy_post;
213        routes.post_completions            = models_routes->proxy_post;
214        routes.post_completions_oai        = models_routes->proxy_post;
215        routes.post_chat_completions       = models_routes->proxy_post;
216        routes.post_control                = models_routes->proxy_post;
217        routes.post_responses_oai          = models_routes->proxy_post;
218        routes.post_transcriptions_oai     = models_routes->proxy_post;
219        routes.post_anthropic_messages     = models_routes->proxy_post;
220        routes.post_anthropic_count_tokens = models_routes->proxy_post;
221        routes.post_infill                 = models_routes->proxy_post;
222        routes.post_embeddings             = models_routes->proxy_post;
223        routes.post_embeddings_oai         = models_routes->proxy_post;
224        routes.post_rerank                 = models_routes->proxy_post;
225        routes.post_tokenize               = models_routes->proxy_post;
226        routes.post_detokenize             = models_routes->proxy_post;
227        routes.post_apply_template         = models_routes->proxy_post;
228        routes.post_chat_completions_tok   = models_routes->proxy_post;
229        routes.post_responses_tok_oai      = models_routes->proxy_post;
230        routes.get_lora_adapters           = models_routes->proxy_get;
231        routes.post_lora_adapters          = models_routes->proxy_post;
232        routes.get_slots                   = models_routes->proxy_get;
233        routes.post_slots                  = models_routes->proxy_post;
234
235        // custom routes for router
236        routes.get_props                   = models_routes->get_router_props;
237        routes.get_models                  = models_routes->get_router_models;
238
239        ctx_http.post("/models",               ex_wrapper(models_routes->post_router_models));
240        ctx_http.post("/models/load",          ex_wrapper(models_routes->post_router_models_load));
241        ctx_http.post("/models/unload",        ex_wrapper(models_routes->post_router_models_unload));
242        ctx_http.get ("/models/sse",           ex_wrapper(models_routes->get_router_models_sse));
243        ctx_http.del ("/models",               ex_wrapper(models_routes->del_router_models));
244    }
245
246    ctx_http.get ("/health",                   ex_wrapper(routes.get_health)); // public endpoint (no API key check)
247    ctx_http.get ("/v1/health",                ex_wrapper(routes.get_health)); // public endpoint (no API key check)
248    ctx_http.get ("/metrics",                  ex_wrapper(routes.get_metrics));
249    ctx_http.get ("/props",                    ex_wrapper(routes.get_props));
250    ctx_http.post("/props",                    ex_wrapper(routes.post_props));
251    ctx_http.get ("/models",                   ex_wrapper(routes.get_models));
252    ctx_http.get ("/v1/models",                ex_wrapper(routes.get_models));
253    ctx_http.post("/completion",               ex_wrapper(routes.post_completions)); // legacy
254    ctx_http.post("/completions",              ex_wrapper(routes.post_completions));
255    ctx_http.post("/v1/completions",           ex_wrapper(routes.post_completions_oai));
256    ctx_http.post("/chat/completions",         ex_wrapper(routes.post_chat_completions));
257    ctx_http.post("/v1/chat/completions",      ex_wrapper(routes.post_chat_completions));
258    ctx_http.post("/v1/chat/completions/control", ex_wrapper(routes.post_control));
259    ctx_http.post("/v1/responses",             ex_wrapper(routes.post_responses_oai));
260    ctx_http.post("/responses",                ex_wrapper(routes.post_responses_oai));
261    ctx_http.post("/v1/audio/transcriptions",  ex_wrapper(routes.post_transcriptions_oai));
262    ctx_http.post("/audio/transcriptions",     ex_wrapper(routes.post_transcriptions_oai));
263    ctx_http.post("/v1/messages",              ex_wrapper(routes.post_anthropic_messages)); // anthropic messages API
264    ctx_http.post("/infill",                   ex_wrapper(routes.post_infill));
265    ctx_http.post("/embedding",                ex_wrapper(routes.post_embeddings)); // legacy
266    ctx_http.post("/embeddings",               ex_wrapper(routes.post_embeddings));
267    ctx_http.post("/v1/embeddings",            ex_wrapper(routes.post_embeddings_oai));
268    ctx_http.post("/rerank",                   ex_wrapper(routes.post_rerank));
269    ctx_http.post("/reranking",                ex_wrapper(routes.post_rerank));
270    ctx_http.post("/v1/rerank",                ex_wrapper(routes.post_rerank));
271    ctx_http.post("/v1/reranking",             ex_wrapper(routes.post_rerank));
272    ctx_http.post("/tokenize",                 ex_wrapper(routes.post_tokenize));
273    ctx_http.post("/detokenize",               ex_wrapper(routes.post_detokenize));
274    ctx_http.post("/apply-template",           ex_wrapper(routes.post_apply_template));
275    // token counting
276    ctx_http.post("/chat/completions/input_tokens",    ex_wrapper(routes.post_chat_completions_tok));
277    ctx_http.post("/v1/chat/completions/input_tokens", ex_wrapper(routes.post_chat_completions_tok));
278    ctx_http.post("/responses/input_tokens",           ex_wrapper(routes.post_responses_tok_oai));
279    ctx_http.post("/v1/responses/input_tokens",        ex_wrapper(routes.post_responses_tok_oai));
280    ctx_http.post("/v1/messages/count_tokens",         ex_wrapper(routes.post_anthropic_count_tokens)); // anthropic token counting
281    // LoRA adapters hotswap
282    ctx_http.get ("/lora-adapters",            ex_wrapper(routes.get_lora_adapters));
283    ctx_http.post("/lora-adapters",            ex_wrapper(routes.post_lora_adapters));
284    // Save & load slots
285    ctx_http.get ("/slots",                    ex_wrapper(routes.get_slots));
286    ctx_http.post("/slots/:id_slot",           ex_wrapper(routes.post_slots));
287
288    // resumable streaming: a child binds the local session factories, the router binds
289    // proxies that resolve the owning child, see server-stream.h
290    server_http_context::handler_t stream_get_h;
291    server_http_context::handler_t streams_lookup_h;
292    server_http_context::handler_t stream_delete_h;
293    if (is_router_server) {
294        stream_get_h     = models_routes->router_stream_get;
295        streams_lookup_h = models_routes->router_streams_lookup;
296        stream_delete_h  = models_routes->router_stream_delete;
297    } else {
298        stream_get_h     = server_stream_make_get_handler();
299        streams_lookup_h = server_stream_make_lookup_handler();
300        stream_delete_h  = server_stream_make_delete_handler();
301    }
302    ctx_http.get ("/v1/stream",                ex_wrapper(stream_get_h));
303    ctx_http.post("/v1/streams/lookup",        ex_wrapper(streams_lookup_h));
304    ctx_http.del ("/v1/stream",                ex_wrapper(stream_delete_h));
305
306    // Google Cloud Platform (Vertex AI) compat
307    ctx_http.register_gcp_compat();
308
309    // return 403 for disabled features
310    server_http_context::handler_t res_403 = [](const server_http_req &) {
311        auto res = std::make_unique<server_http_res>();
312        res->status = 403;
313        res->data = safe_json_to_str({
314            {"error", {
315                {"message", "this feature is disabled"},
316                {"type", "feature_disabled"},
317            }}
318        });
319        return res;
320    };
321
322    if (params.cors_origins == "*" && params.api_keys.empty()) {
323        SRV_WRN("%s", "-----------------\n");
324        SRV_WRN("%s", "CORS is set to allow all origins ('*') and no API key is set\n");
325        SRV_WRN("%s", "this can be a security risk (cross-origin attacks)\n");
326        SRV_WRN("%s", "more info: https://github.com/ggml-org/llama.cpp/pull/25655\n");
327        SRV_WRN("%s", "-----------------\n");
328    }
329
330    // CORS proxy (EXPERIMENTAL, only used by the Web UI for MCP)
331    std::vector<std::string> warn_names;
332    if (is_router_server) {
333        warn_names.push_back("router mode");
334    }
335
336    if (params.ui_mcp_proxy) {
337        ctx_http.get ("/cors-proxy",      ex_wrapper(proxy_handler_get));
338        ctx_http.post("/cors-proxy",      ex_wrapper(proxy_handler_post));
339        warn_names.push_back("MCP proxy (experimental)");
340    } else {
341        ctx_http.get ("/cors-proxy",      ex_wrapper(res_403));
342        ctx_http.post("/cors-proxy",      ex_wrapper(res_403));
343    }
344
345    try {
346        mcp_mgr.start(params);
347    } catch (const std::exception & e) {
348        SRV_ERR("MCP starting failed: %s\n", e.what());
349        return 1;
350    }
351
352    if (!params.server_tools.empty() || !mcp_mgr.empty()) {
353        try {
354            tools.setup(params.server_tools, mcp_mgr, params.server_tools_runtime);
355        } catch (const std::exception & e) {
356            SRV_ERR("tools setup failed: %s\n", e.what());
357            return 1;
358        }
359        ctx_http.get ("/tools",           ex_wrapper(tools.handle_get));
360        ctx_http.post("/tools",           ex_wrapper(tools.handle_post));
361        if (!params.server_tools.empty()) {
362            warn_names.push_back("server tools (experimental)");
363        }
364        if (!params.server_tools_runtime.empty()) {
365            warn_names.push_back("tools runtime (experimental)");
366        }
367        if (!mcp_mgr.empty()) {
368            warn_names.push_back("MCP servers (experimental)");
369        }
370    } else {
371        ctx_http.get ("/tools",           ex_wrapper(res_403));
372        ctx_http.post("/tools",           ex_wrapper(res_403));
373    }
374
375    if (warn_names.size() > 0) {
376        SRV_WRN("%s", "-----------------\n");
377        SRV_WRN("%s", "the following feature(s) are enabled:\n");
378        for (const auto & name : warn_names) {
379            SRV_WRN("    %s\n", name.c_str());
380        }
381        SRV_WRN("%s", "do not expose the server to untrusted environments\n");
382        SRV_WRN("%s", "-----------------\n");
383    }
384
385    //
386    // Handle downloading model
387    //
388
389    if (child.is_child() && child.get_mode() == SERVER_CHILD_MODE_DOWNLOAD) {
390        return child.run_download(params);
391    } else if (!is_router_server && !is_run_by_cli) {
392        // single-model mode (NOT spawned by router)
393        // if this is invoked by CLI, model downloading should be already handled
394        try {
395            common_models_handler_apply(models_handler, params);
396        } catch (const std::exception & e) {
397            SRV_ERR("failed to download model: %s\n", e.what());
398            return 1;
399        }
400    }
401
402    //
403    // Start the server
404    //
405
406    std::function<void()> clean_up;
407
408    if (is_router_server) {
409        SRV_INF("%s", "starting server in router mode. models will be automatically loaded on-demand\n");
410
411        clean_up = [&models_routes, &mcp_mgr]() {
412            SRV_INF("%s: cleaning up before exit...\n", __func__);
413            // stop the session GC first, it finalizes live sessions and wakes pending readers
414            server_stream_session_manager_stop();
415            if (models_routes.has_value()) {
416                models_routes->stopping.store(true); // maybe redundant, but just to be safe
417                models_routes->models.unload_all();
418            }
419            mcp_mgr.shutdown();
420            llama_backend_free();
421        };
422
423        if (!ctx_http.start()) {
424            clean_up();
425            SRV_ERR("%s", "exiting due to HTTP server error\n");
426            return 1;
427        }
428        ctx_http.is_ready.store(true);
429
430        shutdown_handler = [&](int) {
431            if (models_routes.has_value()) {
432                // important to disconnect any SSE clients
433                models_routes->stopping.store(true);
434            }
435            mcp_mgr.shutdown();
436            ctx_http.stop();
437        };
438
439        try {
440            models_routes->models.load_startup_models();
441        } catch (const std::exception & e) {
442            SRV_ERR("failed to load models on startup: %s\n", e.what());
443            ctx_http.stop();
444            if (ctx_http.thread.joinable()) {
445                ctx_http.thread.join();
446            }
447            clean_up();
448            return 1;
449        }
450
451    } else {
452        // setup clean up function, to be called before exit
453        clean_up = [&ctx_http, &ctx_server, &mcp_mgr]() {
454            SRV_INF("%s: cleaning up before exit...\n", __func__);
455            // stop the session GC first, it finalizes live sessions and wakes pending readers
456            server_stream_session_manager_stop();
457            ctx_http.stop();
458            ctx_server.terminate();
459            mcp_mgr.shutdown();
460            llama_backend_free();
461        };
462
463        // start the HTTP server before loading the model to be able to serve /health requests
464        if (!ctx_http.start()) {
465            clean_up();
466            SRV_ERR("%s", "exiting due to HTTP server error\n");
467            return 1;
468        }
469
470        // setup communication child --> router if necessary
471        if (child.is_child()) {
472            ctx_server.set_state_callback([&](server_state state, json payload) {
473                child.notify_to_router(server_state_to_str(state), payload);
474            });
475        }
476
477        if (!ctx_server.load_model(params)) {
478            clean_up();
479            if (ctx_http.thread.joinable()) {
480                ctx_http.thread.join();
481            }
482            SRV_ERR("%s", "exiting due to model loading error\n");
483            return 1;
484        }
485
486        routes.update_meta(ctx_server);
487        ctx_http.is_ready.store(true);
488
489        SRV_INF("%s", "model loaded\n");
490
491        shutdown_handler = [&](int) {
492            mcp_mgr.shutdown();
493            // this will unblock start_loop()
494            ctx_server.terminate();
495        };
496    }
497
498    // register signal handler if not running by CLI
499    if (!is_run_by_cli) {
500#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
501        struct sigaction sigint_action;
502        sigint_action.sa_handler = signal_handler;
503        sigemptyset (&sigint_action.sa_mask);
504        sigint_action.sa_flags = 0;
505        sigaction(SIGINT, &sigint_action, NULL);
506        sigaction(SIGTERM, &sigint_action, NULL);
507#elif defined (_WIN32)
508        auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
509            return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;
510        };
511        SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
512#endif
513    }
514
515    SRV_INF("listening on %s\n", ctx_http.listening_address.c_str());
516
517    // TODO: remove this in the future
518    // check the string to also handle the .sock case
519    if (string_ends_with(ctx_http.listening_address, ":8080")) {
520        SRV_WRN("%s", "NOTICE: server default port will be changed to :9931 in a future release\n");
521        SRV_WRN("%s", "        ref: https://github.com/ggml-org/llama.cpp/pull/26508\n");
522    }
523
524    if (is_router_server) {
525        if (!params.models_preset_hf.empty()) {
526            SRV_WRN(      "NOTE: using preset.ini from HF repo '%s'\n", params.models_preset_hf.c_str());
527            SRV_WRN("%s", "      please only use presets that you can trust! Unknown presets may be unsafe\n");
528        }
529
530        if (ctx_http.thread.joinable()) {
531            ctx_http.thread.join(); // keep the main thread alive
532        }
533
534        // when the HTTP server stops, clean up and exit
535        clean_up();
536    } else {
537        // optionally, notify router server that this instance is ready
538        std::thread monitor_thread;
539        if (child.is_child()) {
540            monitor_thread = child.setup(shutdown_handler);
541            child.notify_to_router(server_state_to_str(SERVER_STATE_READY), routes.get_model_info());
542        }
543
544        // this call blocks the main thread until queue_tasks.terminate() is called
545        ctx_server.start_loop();
546
547        clean_up();
548        if (ctx_http.thread.joinable()) {
549            ctx_http.thread.join();
550        }
551        if (monitor_thread.joinable()) {
552            monitor_thread.join();
553        }
554
555        auto * ll_ctx = ctx_server.get_llama_context();
556        if (ll_ctx != nullptr) {
557            common_memory_breakdown_print(ll_ctx);
558        }
559    }
560
561    return 0;
562}