serve: ONEBIT_WARMUP selects which warm-up request is sent (engine#1539's bisect)

app/serve.cpp

@@ -1649,13 +1649,24 @@ int serve_only(const std::string& why) {
     return kServeOnly;
 }
 
-// The warm-up requests (warm_backend, below): HRX only, and ONEBIT_NO_WARMUP=1 turns them off. serve sends
+// The warm-up requests (warm_backend, below): HRX only, and ONEBIT_NO_WARMUP=1 (or ONEBIT_WARMUP=none)
+// turns them off; ONEBIT_WARMUP picks which of the two, which is how engine#1539's bisect runs. serve sends
 // them itself; --print-backend hands the same requests to Lemonade as the plan's "warmups" ([{path, body}],
 // in order; "warmup" is the first one, for a lemond that reads only that), and Lemonade's onebit recipe
 // sends them once the backend is ready, before it marks the model loaded.
+// ONEBIT_WARMUP selects which warm-up requests are sent (engine#1539's bisect): "" or "both" = the shipped pair,
+// "long" = only the sampling chat request, "short" = only the 16-token /completion prompt, "none" = neither (the
+// spelling ONEBIT_NO_WARMUP already had). The two requests are not interchangeable -- the first samples and the
+// second does not -- so which one a model cannot survive is a question one load can answer instead of a rebuild.
+static std::string warmup_selection() {
+    const char* which = getenv("ONEBIT_WARMUP");
+    return which ? std::string(which) : std::string();
+}
+
 static bool warmup_wanted(const std::string& device) {
     const std::string dev = device.substr(0, device.find(' '));
-    return getenv("ONEBIT_NO_WARMUP") == nullptr && (dev == "hrx" || dev == "hybrid");
+    return getenv("ONEBIT_NO_WARMUP") == nullptr && warmup_selection() != "none" &&
+           (dev == "hrx" || dev == "hybrid");
 }
 
 static const char* const kWarmupPath = "/v1/chat/completions";
@@ -1689,8 +1700,16 @@ static json warmup_short_body() {
 }
 
 static json warmup_requests() {
-    return json::array({ json{ { "path", kWarmupPath }, { "body", warmup_body() } },
-                         json{ { "path", "/completion" }, { "body", warmup_short_body() } } });
+    // An unrecognised value is the shipped pair, not "none": a typo must not quietly disable a warm-up that moves a
+    // measured 2 s off the client's first request. Normalised here so a caller cannot see the two states differently.
+    std::string sel = warmup_selection();
+    if (sel != "long" && sel != "short" && sel != "none") sel = "both";
+    const bool want_chat  = sel == "both" || sel == "long";
+    const bool want_short = sel == "both" || sel == "short";
+    json out = json::array();
+    if (want_chat)  out.push_back(json{ { "path", kWarmupPath },    { "body", warmup_body() } });
+    if (want_short) out.push_back(json{ { "path", "/completion" },  { "body", warmup_short_body() } });
+    return out;
 }
 
 // --print-backend: the llama-server serve would start for this model, as one JSON line on stdout,
@@ -1732,7 +1751,7 @@ int print_backend(const Options& o, const Launch& l) {
     // a chat backend on HRX: the request Lemonade sends before it marks the model loaded (warmup_wanted)
     if (o.role.empty() && warmup_wanted(l.device)) {
         out["warmups"] = warmup_requests();
-        out["warmup"]  = out["warmups"][0];
+        if (!out["warmups"].empty()) out["warmup"] = out["warmups"][0];
     }
     std::printf("%s\n", out.dump().c_str());
     return 0;

docs/serve.md

@@ -333,7 +333,12 @@ What serve did on the request path is the backend's own on this route:
   `"warmups": [{"path", "body"}, ...]` (and `"warmup"`, the first one, for an older lemond), on `hrx`
   chat plans only; Lemonade's onebit recipe sends them in order after `/health` answers and before it
   marks the model loaded. `ONEBIT_NO_WARMUP=1` in lemond's environment leaves them out of the plan, as
-  it turns serve's own warm-up off.
+  it turns serve's own warm-up off. `ONEBIT_WARMUP=long|short|none` selects *which* of the two, in
+  serve's own environment and in the plan alike (`long` = the chat request, `short` = the 16-token
+  completion, `none` = neither; unset or `both` is the shipped pair, and an unrecognised value keeps
+  it). The two are not interchangeable — the first samples, the second does not — so when a model
+  cannot survive one of them, this is the difference between one load per hypothesis and a rebuild
+  (engine#1539).
 
   **Why the 16-token request.** The first batch of 13 to 20 tokens after a load takes ggml-hrx's
   small-batch matmul kernels, which read the weights in a layout ggml-hrx builds on the host the first

tests/print_backend.sh

@@ -29,6 +29,7 @@
 #     --device cpu, --parallel 2, --embedding and a file without the key leave it out,
 #   - an hrx chat plan carries "warmups" (serve's warm-up requests, in order: the chat request, then a
 #     16-token /completion prompt (" a" x 16, n_predict 0)) and "warmup" (the first one, for an older lemond),
+#   - ONEBIT_WARMUP=long|short|none picks which warm-up request(s) the plan carries (engine#1539's bisect)
 #   - an hrx chat plan carries "warmup" (serve's warm-up request, for Lemonade to send before the
 #     model counts as loaded); ONEBIT_NO_WARMUP=1, cpu and role servers have none.
 #
@@ -180,6 +181,27 @@ check "hrx plan: warmups = [the chat request, a 16-token /completion], warmup =
 out=$(ONEBIT_NO_WARMUP=1 "$bin" serve -m "$scratch/small.gguf" --device hrx --llama-server "$here/fake_backend.py" \
       --no-mem-check --port 43132 --print-backend 2>"$scratch/err"); rc=$?
 check "hrx plan: ONEBIT_NO_WARMUP=1, no warmup (exit $rc)" '[ "$rc" = 0 ] && [ "$(get "\"warmup\" in d or \"warmups\" in d")" = False ]'
+# engine#1539: the two warm-ups are not interchangeable (the first samples, the second does not), so which one a
+# model cannot survive is now a selector rather than a rebuild. ONEBIT_WARMUP picks them; "none" is the spelling
+# ONEBIT_NO_WARMUP already had, and an unknown value keeps the shipped pair rather than silently sending none.
+for _sel in long short; do
+    out=$(ONEBIT_WARMUP=$_sel "$bin" serve -m "$scratch/small.gguf" --device hrx --llama-server "$here/fake_backend.py" \
+          --no-mem-check --port 43133 --print-backend 2>"$scratch/err"); rc=$?
+    if [ "$_sel" = long ]; then _want=/v1/chat/completions; else _want=/completion; fi
+    # A join, not print(list): python renders a one-element list as ['/completion'], quotes included, and comparing
+    # against a bare [$_want] is a test that fails on a correct plan (found by running it).
+    _got=$(get "' '.join(w[\"path\"] for w in d[\"warmups\"])")
+    _mirrors=$(get "d[\"warmup\"] == d[\"warmups\"][0]")
+    check "hrx plan: ONEBIT_WARMUP=$_sel sends only $_want (exit $rc)" \
+        '[ "$rc" = 0 ] && [ "$_got" = "$_want" ] && [ "$_mirrors" = True ]'
+done
+out=$(ONEBIT_WARMUP=none "$bin" serve -m "$scratch/small.gguf" --device hrx --llama-server "$here/fake_backend.py" \
+      --no-mem-check --port 43134 --print-backend 2>"$scratch/err"); rc=$?
+check "hrx plan: ONEBIT_WARMUP=none, no warmup (exit $rc)" '[ "$rc" = 0 ] && [ "$(get "\"warmup\" in d or \"warmups\" in d")" = False ]'
+out=$(ONEBIT_WARMUP=bogus "$bin" serve -m "$scratch/small.gguf" --device hrx --llama-server "$here/fake_backend.py" \
+      --no-mem-check --port 43135 --print-backend 2>"$scratch/err"); rc=$?
+check "hrx plan: an unknown ONEBIT_WARMUP keeps the shipped pair (exit $rc)" \
+    '[ "$rc" = 0 ] && [ "$(get "len(d[\"warmups\"])")" = 2 ]'
 out=$(mplan small cpu); rc=$?
 check "cpu plan: no warmup (exit $rc)" '[ "$rc" = 0 ] && [ "$(get "\"warmup\" in d")" = False ]'