server: POST /lora-adapters/snapshot writes a mind's shadow to an immutable GGUF between steps, and /learn reports the shadow's committed step

tools/server/server-context.cpp

@@ -903,6 +903,8 @@ struct server_context_impl {
     // LEARNING IN THE SERVING CONTEXT (SERVER_TASK_TYPE_LEARN): the learning rate the next step reads
     // (each mind's shadow has its own optimizer in the context: llama_opt_init selects it)
     float   learn_lr      = 1e-4f;
+    // each shadow's committed optimizer steps on this server (a snapshot records where she was)
+    std::map<const llama_adapter_lora *, int64_t> learn_steps;
     // LEARNING BARELY TOUCHES THE FLOW: while any other slot is working, learning may take at most
     // LEARN_DUTY of the loop's time (the step after a learn of t ms waits t*(1/LEARN_DUTY - 1) ms);
     // with every other slot idle it runs at once. Back to back, learns took the 4B's other slot from
@@ -2207,6 +2209,7 @@ struct server_context_impl {
         slot.prompt.checkpoints.remove_if([&](const common_prompt_checkpoint & c) { return c.n_tokens > n_keep; });
         learn_outcome out;
         out.ok       = ok;
+        out.step     = ok ? ++learn_steps[loras[L.lora_id].ptr] : 0;
         out.failure  = ok ? "" : llama_opt_failure(ctx_tgt);
         out.loss     = loss;
         out.n_labels = (int32_t) toks.size() - 1;
@@ -3166,6 +3169,44 @@ struct server_context_impl {
                     res->outcome = learn_on_slot(*slot, L, t0);
                     queue_results.send(std::move(res));
                 } break;
+            case SERVER_TASK_TYPE_SNAPSHOT_LORA:
+                {
+                    // Between steps on the loop, so the file is one whole step's weights, never a
+                    // half-applied update. Written beside its name and renamed into place, and
+                    // never over an existing file: a snapshot is immutable.
+                    auto & loras = params_base.lora_adapters;
+                    const int32_t id = task.snapshot_lora;
+                    if (id < 0 || (size_t) id >= loras.size() || loras[id].ptr == nullptr) {
+                        send_error(task, "no loaded LoRA adapter with id " + std::to_string(id), ERROR_TYPE_INVALID_REQUEST);
+                        break;
+                    }
+                    std::error_code ec;
+                    if (std::filesystem::exists(task.snapshot_path, ec)) {
+                        send_error(task, "a snapshot is immutable: " + task.snapshot_path + " already exists", ERROR_TYPE_INVALID_REQUEST);
+                        break;
+                    }
+                    const std::string partial = task.snapshot_path + ".partial";
+                    if (llama_adapter_lora_save(loras[id].ptr, partial.c_str()) != 0) {
+                        std::filesystem::remove(partial, ec);
+                        send_error(task, "could not write the snapshot of adapter " + std::to_string(id), ERROR_TYPE_SERVER);
+                        break;
+                    }
+                    std::filesystem::rename(partial, task.snapshot_path, ec);
+                    if (ec) {
+                        std::filesystem::remove(partial, ec);
+                        send_error(task, "could not place the snapshot at " + task.snapshot_path, ERROR_TYPE_SERVER);
+                        break;
+                    }
+                    auto res = std::make_unique<server_task_result_lora_snapshot>();
+                    res->id      = task.id;
+                    res->lora_id = id;
+                    res->path    = task.snapshot_path;
+                    res->bytes   = (int64_t) std::filesystem::file_size(task.snapshot_path, ec);
+                    const auto steps = learn_steps.find(loras[id].ptr);
+                    res->steps   = steps == learn_steps.end() ? 0 : steps->second;
+                    SRV_INF("snapshot of lora adapter id=%d at %lld steps: %s\n", id, (long long) res->steps, res->path.c_str());
+                    queue_results.send(std::move(res));
+                } break;
             case SERVER_TASK_TYPE_UNLOAD_LORA:
                 {
                     auto & loras = params_base.lora_adapters;
@@ -3195,6 +3236,7 @@ struct server_context_impl {
                     // a shadow that learned holds an optimizer referencing its tensors: release it
                     // first (its moments go; a snapshot is the durable form of what it learned)
                     llama_opt_release(ctx_tgt, ptr);
+                    learn_steps.erase(ptr);
                     // An idle slot that last ran with it keeps KV computed WITH the adapter;
                     // once the adapter is gone its list would read as the base's, so that
                     // cache must not be reused.
@@ -6112,6 +6154,37 @@ void server_routes::init_routes() {
         return res;
     };
 
+    this->post_lora_adapters_snapshot = [this](const server_http_req & req) {
+        auto res = create_response();
+        const json body = json::parse(req.body);
+        const std::string name = json_value(body, "name", std::string());
+        const std::string path = confine_out(params.train_dir, name);
+        if (!body.contains("id") || !body.at("id").is_number_integer() || params.train_dir.empty() || path.empty()) {
+            res->error(format_error_response("\"id\" (an integer adapter id) and \"name\" (a bare file name ending in .gguf "
+                                             "inside the server's --train-dir) are required", ERROR_TYPE_INVALID_REQUEST));
+            return res;
+        }
+        auto & rd = res->rd;
+        {
+            server_task task(SERVER_TASK_TYPE_SNAPSHOT_LORA);
+            task.id            = rd.get_new_id();
+            task.snapshot_lora = body.at("id").get<int32_t>();
+            task.snapshot_path = path;
+            rd.post_task(std::move(task));
+        }
+        auto result = rd.next(req.should_stop);
+        if (!result) {
+            GGML_ASSERT(req.should_stop());
+            return res;
+        }
+        if (result->is_error()) {
+            res->error(result->to_json());
+            return res;
+        }
+        res->ok(result->to_json());
+        return res;
+    };
+
     this->post_lora_adapters_unload = [this](const server_http_req & req) {
         auto res = create_response();
         const json body = json::parse(req.body);

tools/server/server-context.h

@@ -169,6 +169,7 @@ struct server_routes {
     server_http_context::handler_t post_lora_adapters;
     server_http_context::handler_t post_lora_adapters_load;
     server_http_context::handler_t post_lora_adapters_unload;
+    server_http_context::handler_t post_lora_adapters_snapshot;
     server_http_context::handler_t post_learn;
     server_http_context::handler_t get_train;
     server_http_context::handler_t post_train;

tools/server/server-task.cpp

@@ -1727,6 +1727,9 @@ json server_task_result_lora_change::to_json() {
 
 json learn_outcome::to_json() const {
     json res = { { "success", ok }, { "loss", loss }, { "n_labels", n_labels }, { "t_ms", t_ms } };
+    if (ok) {
+        res["step"] = step;
+    }
     if (deferred) {
         res["deferred"] = true;
         res["retry_after_ms"] = retry_after_ms;

tools/server/server-task.h

@@ -60,6 +60,7 @@ enum server_task_type {
     SERVER_TASK_TYPE_SET_LORA,
     SERVER_TASK_TYPE_LOAD_LORA,
     SERVER_TASK_TYPE_UNLOAD_LORA,
+    SERVER_TASK_TYPE_SNAPSHOT_LORA,
     SERVER_TASK_TYPE_LEARN,
 };
 
@@ -101,6 +102,8 @@ struct learn_outcome {
     double      loss     = 0.0;
     int32_t     n_labels = 0;
     double      t_ms     = 0.0;
+    // the shadow's committed optimizer steps after this one (0 when it did not step)
+    int64_t     step     = 0;
     json to_json() const;
 };
 
@@ -264,6 +267,10 @@ struct server_task {
     // used by SERVER_TASK_TYPE_UNLOAD_LORA: the adapter id to retire
     int32_t unload_lora = -1;
 
+    // used by SERVER_TASK_TYPE_SNAPSHOT_LORA: the adapter id to write, and where (confined to --train-dir)
+    int32_t     snapshot_lora = -1;
+    std::string snapshot_path;
+
     // used by SERVER_TASK_TYPE_LEARN
     learn_request learn;
 
@@ -667,6 +674,17 @@ struct server_task_result_lora_change : server_task_result {
     virtual json to_json() override;
 };
 
+// A shadow written to an immutable file between steps: what she had learned at `steps`.
+struct server_task_result_lora_snapshot : server_task_result {
+    int32_t     lora_id = -1;
+    std::string path;
+    int64_t     bytes   = 0;
+    int64_t     steps   = 0;
+    virtual json to_json() override {
+        return { { "success", true }, { "id", lora_id }, { "path", path }, { "bytes", bytes }, { "steps", steps } };
+    }
+};
+
 struct server_task_result_learn : server_task_result {
     learn_outcome outcome;
     virtual json to_json() override { return outcome.to_json(); }

tools/server/server.cpp

@@ -219,6 +219,7 @@ int llama_server(common_params & params, int argc, char ** argv) {
         routes.post_lora_adapters          = models_routes->proxy_post;
         routes.post_lora_adapters_load     = models_routes->proxy_post;
         routes.post_lora_adapters_unload   = models_routes->proxy_post;
+        routes.post_lora_adapters_snapshot = models_routes->proxy_post;
         routes.post_learn                  = models_routes->proxy_post;
         routes.get_slots                   = models_routes->proxy_get;
         routes.post_slots                  = models_routes->proxy_post;
@@ -274,6 +275,7 @@ int llama_server(common_params & params, int argc, char ** argv) {
     ctx_http.post("/lora-adapters",            ex_wrapper(routes.post_lora_adapters));
     ctx_http.post("/lora-adapters/load",       ex_wrapper(routes.post_lora_adapters_load));
     ctx_http.post("/lora-adapters/unload",     ex_wrapper(routes.post_lora_adapters_unload));
+    ctx_http.post("/lora-adapters/snapshot",   ex_wrapper(routes.post_lora_adapters_snapshot));
     ctx_http.post("/learn",                    ex_wrapper(routes.post_learn));
     ctx_http.get ("/train",                    ex_wrapper(routes.get_train));
     ctx_http.post("/train",                    ex_wrapper(routes.post_train));

tools/server/tests/learn-every-format.py

@@ -37,6 +37,29 @@ def raw(path, body):
     print(json.dumps({"unbound_retired": retired, "bound_retired": bound, "next_learn": following["learn"]}))
     sys.exit(0)
 
+# A snapshot is what she had learned, immutable: it carries the shadow's weights at its committed step
+# (fresh server; writes snap-test.gguf into --train-dir, so remove it before a rerun).
+if "--snapshot" in sys.argv:
+    import urllib.error
+    turn = {"prompt": "The capital of Kansas is", "id_slot": 0, "n_predict": 1, "temperature": 0, "learn": learn}
+    steps = [json.loads(raw("/completion", turn))["learn"] for _ in range(2)]
+    assert [st["success"] for st in steps] == [True, True] and [st["step"] for st in steps] == [1, 2], steps
+    snap = json.loads(raw("/lora-adapters/snapshot", {"id": sid, "name": "snap-test.gguf"}))
+    assert snap["success"] and snap["steps"] == 2 and snap["bytes"] > 0, snap
+    try:
+        raw("/lora-adapters/snapshot", {"id": sid, "name": "snap-test.gguf"})
+        raise AssertionError("a snapshot was overwritten")
+    except urllib.error.HTTPError as error:
+        detail = error.read().decode()
+        assert error.code == 400 and "immutable" in detail, (error.code, detail)
+    third = json.loads(raw("/completion", turn))["learn"]
+    # a learn's loss is its forward pass before its update: the snapshot, loaded, sees what the shadow saw
+    loaded = json.loads(raw("/lora-adapters/load", {"name": "snap-test.gguf"}))["id"]
+    again = json.loads(raw("/completion", dict(turn, learn=dict(learn, lora_id=loaded))))["learn"]
+    assert again["success"] and abs(again["loss"] - third["loss"]) < 1e-5, (third, again)
+    print(json.dumps({"steps": steps, "snapshot": snap, "third": third, "from_snapshot": again}))
+    sys.exit(0)
+
 # Two minds on one serving context: each shadow keeps its own optimizer while the other learns in
 # between (needs shadow-05.gguf and a copy of it, shadow-06.gguf, in --train-dir).
 if "--two-minds" in sys.argv: