server: reconfigure grows or shrinks the contexts live, on the s4 integration branch (S2 port of #104)

.github/workflows/server.yml

@@ -121,11 +121,23 @@ jobs:
 
       - name: Python setup
         id: setup_python
+        # also after a failed step above, so the reconfigure tests below still run (see their note)
+        if: ${{ !cancelled() }}
         uses: actions/setup-python@v6
         with:
           python-version: '3.11'
           pip-install: -r tools/server/tests/requirements.txt
 
+      # Continuum fork: the live reconfigure's validator (SERVING-ENGINE.md, S2), run on its own while
+      # pre-existing residency failures (the dream step above, unit/test_lora.py) stop the suite under -x
+      # before it; drop this step once the suite reaches test_reconfigure.py again
+      - name: Reconfigure tests
+        if: ${{ !cancelled() }}
+        run: |
+          cd tools/server/tests
+          ./tests.sh -v -x unit/test_reconfigure.py
+          LLAMA_ARG_BACKEND_SAMPLING=1 ./tests.sh -v -x unit/test_reconfigure.py
+
       - name: Tests
         id: server_integration_tests
         run: |
@@ -205,6 +217,14 @@ jobs:
           python-version: '3.11'
           pip-install: -r tools/server/tests/requirements.txt
 
+      # Continuum fork: the live reconfigure's validator, on its own (see the ubuntu job's note)
+      - name: Reconfigure tests
+        shell: bash
+        run: |
+          cd tools/server/tests
+          export PYTHONIOENCODING=":replace"
+          ./tests.sh -v -x unit/test_reconfigure.py
+
       - name: Tests
         id: server_integration_tests
         shell: bash

common/common.cpp

@@ -1278,11 +1278,13 @@ struct common_init_result::impl {
     // note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top
 
     common_threadpools threadpools;
+    bool threadpools_made = false; // init() runs once; a context created later is attached to them
 
     llama_model_ptr   model;
     llama_context_ptr context;
 
     std::vector<llama_adapter_lora_ptr> lora;
+    bool adapters_failed = false; // an adapter did not load: no context is created on this model
 
     std::vector<common_sampler_ptr> samplers;
     std::vector<llama_sampler_seq_config> samplers_seq_config;
@@ -1291,9 +1293,10 @@ struct common_init_result::impl {
 common_init_result::common_init_result(common_params & params, bool model_only) :
     pimpl(new impl{}) {
     auto mparams = common_model_params_to_llama(params);
-    auto cparams = common_context_params_to_llama(params);
 
     if (params.fit_params) {
+        auto cparams = common_context_params_to_llama(params);
+
         COM_TRC("%s", "fitting params to device memory ...\n");
         COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n");
 
@@ -1326,6 +1329,10 @@ common_init_result::common_init_result(common_params & params, bool model_only)
             params.fit_params_min_ctx,
             has_draft || spec_mtp ? &extra : nullptr,
             params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
+
+        // Continuum fork: the context size fit chose becomes the params' own (fit changes only the
+        // context's n_ctx, common/fit.cpp), so init_context() creates the context it fitted
+        params.n_ctx = cparams.n_ctx;
     }
 
     llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams);
@@ -1352,6 +1359,7 @@ common_init_result::common_init_result(common_params & params, bool model_only)
         lora.reset(llama_adapter_lora_init(model, la.path.c_str()));
         if (lora == nullptr) {
             COM_ERR("failed to load lora adapter '%s'\n", la.path.c_str());
+            pimpl->adapters_failed = true;
             return;
         }
 
@@ -1387,8 +1395,23 @@ common_init_result::common_init_result(common_params & params, bool model_only)
                 params.sampling.logit_bias.end(),
                 params.sampling.logit_bias_eog.begin(), params.sampling.logit_bias_eog.end());
     }
+}
+
+bool common_init_result::init_context(common_params & params, bool warmup) {
+    GGML_ASSERT(pimpl->model && "init_context needs a loaded model");
+    GGML_ASSERT(!pimpl->context && "free_context() before creating another context");
+
+    if (pimpl->adapters_failed) {
+        COM_ERR("%s", "an adapter failed to load, no context is created on this model\n");
+        return false;
+    }
+
+    llama_model * model = pimpl->model.get();
+    auto cparams = common_context_params_to_llama(params);
 
     // init the backend samplers as part of the context creation
+    pimpl->samplers.clear();
+    pimpl->samplers_seq_config.clear();
     pimpl->samplers.resize(cparams.n_seq_max);
     pimpl->samplers_seq_config.resize(cparams.n_seq_max);
 
@@ -1405,59 +1428,21 @@ common_init_result::common_init_result(common_params & params, bool model_only)
     llama_context * lctx = llama_init_from_model(model, cparams);
     if (lctx == NULL) {
         COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str());
-        return;
+        return false;
     }
 
     pimpl->context.reset(lctx);
 
     set_process_priority(params.cpuparams.priority);
 
-    pimpl->threadpools.init(lctx, params);
-}
-
-llama_model * common_init_result::model() {
-    return pimpl->model.get();
-}
-
-llama_context * common_init_result::context() {
-    return pimpl->context.get();
-}
-
-common_sampler * common_init_result::sampler(llama_seq_id seq_id) {
-    if (seq_id < 0 || seq_id >= (int) pimpl->samplers.size()) {
-        return nullptr;
-    }
-    return pimpl->samplers[seq_id].get();
-}
-
-void common_init_result::reset_samplers() {
-    for (int i = 0; i < (int) pimpl->samplers.size(); ++i) {
-        llama_sampler_reset(common_sampler_get(pimpl->samplers[i].get()));
-    }
-}
-
-std::vector<llama_adapter_lora_ptr> & common_init_result::lora() {
-    return pimpl->lora;
-}
-
-common_init_result_ptr common_init_from_params(common_params & params, bool model_only) {
-    common_init_result_ptr res(new common_init_result(params, model_only));
-
-    llama_model * model = res->model();
-    if (model == NULL) {
-        COM_ERR("failed to load model '%s'\n", params.model.path.c_str());
-        return res;
-    }
-
-    if (model_only) {
-        return res;
+    if (pimpl->threadpools_made) {
+        pimpl->threadpools.attach(lctx);
+    } else {
+        pimpl->threadpools.init(lctx, params);
+        pimpl->threadpools_made = true;
     }
 
-    llama_context * lctx = res->context();
-    if (lctx == NULL) {
-        COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str());
-        return res;
-    }
+    // the context's own setup below; a failure in it leaves the context created, as it always has
 
     const llama_vocab * vocab = llama_model_get_vocab(model);
 
@@ -1472,7 +1457,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
 
         const auto cvec = common_control_vector_load(params.control_vectors);
         if (cvec.n_embd == -1) {
-            return res;
+            return true;
         }
 
         int err = llama_set_adapter_cvec(
@@ -1483,7 +1468,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
                 params.control_vector_layer_start,
                 params.control_vector_layer_end);
         if (err) {
-            return res;
+            return true;
         }
     }
 
@@ -1507,15 +1492,15 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
         }
 
         if (!ok) {
-            return res;
+            return true;
         }
     }
 
     if (!params.lora_init_without_apply) {
         common_set_adapter_lora(lctx, params.lora_adapters);
     }
 
-    if (params.warmup) {
+    if (warmup) {
         COM_TRC("%s", "warming up the model with an empty run - please wait ... (--no-warmup to disable)\n");
 
         std::vector<llama_token> tmp;
@@ -1550,7 +1535,67 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
         llama_perf_context_reset(lctx);
 
         // reset samplers to reset RNG state after warmup to the seeded state
-        res->reset_samplers();
+        reset_samplers();
+    }
+
+    return true;
+}
+
+void common_init_result::free_context() {
+    pimpl->context.reset();
+    // the backend sampler configs were the freed context's
+    pimpl->samplers.clear();
+    pimpl->samplers_seq_config.clear();
+}
+
+llama_model * common_init_result::model() {
+    return pimpl->model.get();
+}
+
+llama_context * common_init_result::context() {
+    return pimpl->context.get();
+}
+
+common_sampler * common_init_result::sampler(llama_seq_id seq_id) {
+    if (seq_id < 0 || seq_id >= (int) pimpl->samplers.size()) {
+        return nullptr;
+    }
+    return pimpl->samplers[seq_id].get();
+}
+
+void common_init_result::reset_samplers() {
+    for (int i = 0; i < (int) pimpl->samplers.size(); ++i) {
+        llama_sampler_reset(common_sampler_get(pimpl->samplers[i].get()));
+    }
+}
+
+std::vector<llama_adapter_lora_ptr> & common_init_result::lora() {
+    return pimpl->lora;
+}
+
+common_init_result_ptr common_init_from_params(common_params & params, bool model_only) {
+    common_init_result_ptr res(new common_init_result(params, model_only));
+
+    if (res->model() == NULL) {
+        COM_ERR("failed to load model '%s'\n", params.model.path.c_str());
+        return res;
+    }
+
+    if (model_only) {
+        return res;
+    }
+
+    // on failure the context stays null, which is what callers check
+    res->init_context(params, params.warmup);
+
+    return res;
+}
+
+common_init_result_ptr common_load_from_params(common_params & params) {
+    common_init_result_ptr res(new common_init_result(params, false));
+
+    if (res->model() == NULL) {
+        COM_ERR("failed to load model '%s'\n", params.model.path.c_str());
     }
 
     return res;
@@ -1833,6 +1878,13 @@ void common_threadpools::init(llama_context * ctx, const common_params & params)
     llama_attach_threadpool(ctx, threadpool, threadpool_batch);
 }
 
+void common_threadpools::attach(llama_context * ctx) const {
+    if (threadpool == nullptr) {
+        return; // init() made none (no CPU backend, or the pools failed): the context keeps its own
+    }
+    llama_attach_threadpool(ctx, threadpool, threadpool_batch);
+}
+
 //
 // Batch utils
 //

common/common.h

@@ -921,9 +921,18 @@ struct common_sampler;
 
 // note: defines the model, context, samplers, ets. lifetimes
 struct common_init_result {
+    // loads the model and, unless model_only, its adapters; the context is init_context()'s
     common_init_result(common_params & params, bool model_only = false);
     ~common_init_result();
 
+    // Continuum fork: the context is its own lifetime, apart from the model's (SERVING-ENGINE.md,
+    // S2). init_context() creates it on the loaded model with params' geometry and runs its setup
+    // (control vectors, adapters, optionally the warmup); free_context() frees it and what is bound
+    // to it. The model, its mmap, its adapters and the CPU threadpools stay, so a server changes its
+    // geometry without loading the weights again. Returns false if no context could be created.
+    bool init_context(common_params & params, bool warmup);
+    void free_context();
+
     llama_model * model();
     llama_context * context();
 
@@ -939,8 +948,13 @@ struct common_init_result {
 
 using common_init_result_ptr = std::unique_ptr<common_init_result>;
 
+// the model and its context (model_only: the model alone)
 common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false);
 
+// Continuum fork: the model and its adapters with no context yet; the caller creates it with
+// init_context(), and may free and create it again with another geometry
+common_init_result_ptr common_load_from_params(common_params & params);
+
 struct llama_model_params   common_model_params_to_llama  (      common_params & params);
 struct llama_context_params common_context_params_to_llama(const common_params & params);
 
@@ -968,6 +982,9 @@ struct common_threadpools {
 
     void init(llama_context * ctx, const common_params & params);
 
+    // Continuum fork: the pools already made, attached to a context created after them
+    void attach(llama_context * ctx) const;
+
 private:
     ggml_threadpool * threadpool       = nullptr;
     ggml_threadpool * threadpool_batch = nullptr;