Skip to content

Commit 721354f

Browse files
authored
server: (router) move model downloading to dedicated process (ggml-org#24834)
* server: real-time model load progress tracking via /models/sse * update docs * server: move model download to child process * rm unused * fix most problems * clean up * nit fixes * fix test case * do not detact() thread * shorter MODEL_DOWNLOAD_TIMEOUT in test * throttle
1 parent 6ee0f65 commit 721354f

9 files changed

Lines changed: 310 additions & 150 deletions

File tree

common/arg.cpp

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -396,7 +396,7 @@ static bool parse_bool_value(const std::string & value) {
396396
// CLI argument parsing functions
397397
//
398398

399-
bool common_params_handle_models(common_params & params, llama_example curr_ex) {
399+
bool common_params_handle_models(common_params & params, llama_example curr_ex, common_download_callback * callback) {
400400
const bool spec_type_draft_mtp = std::find(params.speculative.types.begin(),
401401
params.speculative.types.end(),
402402
COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
@@ -408,6 +408,10 @@ bool common_params_handle_models(common_params & params, llama_example curr_ex)
408408
opts.download_mtp = spec_type_draft_mtp;
409409
opts.download_mmproj = !params.no_mmproj && params.mmproj.path.empty() && params.mmproj.url.empty();
410410

411+
if (callback) {
412+
opts.callback = callback;
413+
}
414+
411415
// sub-models (draft, mmproj, vocoder) are explicitly specified by the user,
412416
// so we should not auto-discover mtp/mmproj siblings for them
413417
common_download_opts sub_opts = opts;
@@ -584,8 +588,11 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
584588
throw std::invalid_argument("error: --prompt-cache-all not supported in interactive mode yet\n");
585589
}
586590

587-
// export_graph_ops loads only metadata
588-
const bool skip_model_download = ctx_arg.ex == LLAMA_EXAMPLE_EXPORT_GRAPH_OPS;
591+
const bool skip_model_download =
592+
// server will call common_params_handle_models() later, so we skip it here
593+
ctx_arg.ex == LLAMA_EXAMPLE_SERVER ||
594+
// export_graph_ops loads only metadata
595+
ctx_arg.ex == LLAMA_EXAMPLE_EXPORT_GRAPH_OPS;
589596

590597
if (!skip_model_download) {
591598
// handle model and download
@@ -594,7 +601,6 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
594601
// model is required (except for server)
595602
// TODO @ngxson : maybe show a list of available models in CLI in this case
596603
if (params.model.path.empty()
597-
&& ctx_arg.ex != LLAMA_EXAMPLE_SERVER
598604
&& !params.usage
599605
&& !params.completion) {
600606
throw std::invalid_argument("error: --model is required\n");

common/arg.h

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#pragma once
22

33
#include "common.h"
4+
#include "download.h"
45

56
#include <set>
67
#include <map>
@@ -133,7 +134,10 @@ void common_params_add_preset_options(std::vector<common_arg> & args);
133134
// return true if the model is ready to use
134135
// throw an exception if there is an error that prevents the model from being used (e.g. network error, model not found, etc)
135136
// if params.skip_download is true, no downloads will be attempted. return false if the model is invalid or missing (e.g. ETag check failed)
136-
bool common_params_handle_models(common_params & params, llama_example curr_ex);
137+
bool common_params_handle_models(
138+
common_params & params,
139+
llama_example curr_ex,
140+
common_download_callback * callback = nullptr);
137141

138142
// initialize argument parser context - used by test-arg-parser and preset
139143
common_params_context common_params_parser_init(common_params & params, llama_example ex, void(*print_usage)(int, char **) = nullptr);

tools/server/README-dev.md

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -204,9 +204,9 @@ Instead of building everything from the ground up (like what most AI agents will
204204

205205
The flow for downloading a new model:
206206
- POST request comes in --> `post_router_models` --> validation
207-
- `server_models::download()` is called
208-
- Sets up a new thread `inst.th` and runs the download inside
209-
- If a stop request comes in, set `stop_download` to `true`
207+
- A new `llama-server` subprocess will be spawned with special `SERVER_CHILD_MODE_DOWNLOAD`
208+
- Child process runs the download and report status back to router via stdin/out
209+
- If a stop request comes in, the router asks the child process to stop (same mechanism as running a model in child process)
210210
- Otherwise, upon completion, we call `load_models()` to refresh the list of models
211211

212212
### Notable Related PRs

tools/server/server-context.cpp

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -931,6 +931,8 @@ struct server_context_impl {
931931

932932
bool sleeping = false;
933933

934+
int64_t t_last_load_progress_ms = 0;
935+
934936
void destroy() {
935937
spec.reset();
936938
ctx_dft.reset();
@@ -1244,6 +1246,10 @@ struct server_context_impl {
12441246
}
12451247

12461248
if (has_mmproj) {
1249+
if (callback_state) {
1250+
callback_state(SERVER_STATE_LOADING, {{"stage", "mmproj_model"}});
1251+
}
1252+
12471253
if (!is_resume) {
12481254
mtmd_helper_log_set(common_log_default_callback, nullptr);
12491255
}

tools/server/server-context.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,14 +53,15 @@ struct server_context_meta {
5353
};
5454

5555
enum server_state {
56-
// SERVER_STATE_DOWNLOADING,
56+
SERVER_STATE_DOWNLOADING,
5757
SERVER_STATE_LOADING,
5858
SERVER_STATE_READY,
5959
SERVER_STATE_SLEEPING,
6060
};
6161

6262
static std::string server_state_to_str(server_state state) {
6363
switch (state) {
64+
case SERVER_STATE_DOWNLOADING: return "downloading";
6465
case SERVER_STATE_LOADING: return "loading";
6566
case SERVER_STATE_READY: return "ready";
6667
case SERVER_STATE_SLEEPING: return "sleeping";
@@ -69,6 +70,7 @@ static std::string server_state_to_str(server_state state) {
6970
}
7071

7172
static server_state server_state_from_str(const std::string & str) {
73+
if (str == "downloading") return SERVER_STATE_DOWNLOADING;
7274
if (str == "loading") return SERVER_STATE_LOADING;
7375
if (str == "ready") return SERVER_STATE_READY;
7476
if (str == "sleeping") return SERVER_STATE_SLEEPING;

0 commit comments

Comments
 (0)