server : put the preemption record inside the completed response of a streamed /v1/responses

A non-streamed /v1/responses carries preempt in the response object, next to usage. The streamed one wrote it on the SSE data beside the response object, so a client that keeps the response of the response.completed event, which is the object the OpenAI SDK hands back, never saw it. It now sits in the same place either way, and unconditionally, as the non-streamed body already did.

Author: danielhanchen

tools/server/README.md

@@ -677,7 +677,7 @@ These words will not be included in the completion, so make sure to add them to
 - `tokens_cached`: Number of tokens from the prompt which could be re-used from previous completion
 - `tokens_evaluated`: Number of tokens evaluated in total from the prompt
 - `truncated`: Boolean indicating if the context size was exceeded during generation, i.e. the number of tokens provided in the prompt (`tokens_evaluated`) plus tokens generated (`tokens predicted`) exceeded the context size (`n_ctx`)
-- `preempt`: How the request was served while the unified KV cache was full (see `--preempt-ram`). `parks` is how often the request was parked to make room for another, and `recomputes` is how many of those parks dropped the sequence's cells because `--preempt-ram` was spent, so that the resume re-prefilled its tokens instead of restoring the bytes that were saved. A re-prefilled sequence continues from the same tokens, but its numerics are not guaranteed identical to the sequence that left, `LLAMA_EXACT_CONCURRENCY` included: raise `--preempt-ram` until `recomputes` stays 0 where that matters. Both fields are present in the final response of a streamed completion as well.
+- `preempt`: How the request was served while the unified KV cache was full (see `--preempt-ram`). `parks` is how often the request was parked to make room for another, and `recomputes` is how many of those parks dropped the sequence's cells because `--preempt-ram` was spent, so that the resume re-prefilled its tokens instead of restoring the bytes that were saved. A re-prefilled sequence continues from the same tokens, but its numerics are not guaranteed identical to the sequence that left, `LLAMA_EXACT_CONCURRENCY` included: raise `--preempt-ram` until `recomputes` stays 0 where that matters. Both fields are present in the final response of a streamed completion as well, and on the OpenAI-compatible endpoints: the final chunk of a streamed `/v1/chat/completions`, the `message_delta` event of `/v1/messages`, and the response object of `/v1/responses` streamed or not, which in a stream is the `response` of the `response.completed` event.
 
 While a request is streaming, the server sends SSE comment lines that a client reading raw lines can act on and every SSE event consumer ignores:
 

tools/server/server-task.cpp

@@ -714,14 +714,15 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
                     {"output_tokens", n_decoded},
                     {"total_tokens",  n_decoded + n_prompt_tokens},
                     {"input_tokens_details", json { {"cached_tokens", n_prompt_tokens_cache} }},
-                }}
+                }},
+                // [TAG_PREEMPT] inside the response object, where the non-streaming body carries it: that object is what a client keeps from the stream
+                {"preempt",    preempt_to_json()},
             }},
         }}
     });
 
     if (stats.is_set()) {
         server_sent_events.back().at("data")["timings"] = stats.to_json();
-        server_sent_events.back().at("data")["preempt"] = preempt_to_json();
     }
 
     return server_sent_events;

tools/server/tests/unit/test_preempt.py

@@ -940,6 +940,42 @@ def test_a_swap_park_is_not_reported_as_a_recompute():
     assert _metrics()["preempt_recompute_total"] == 0
 
 
+def _stream_responses(n_predict: int, prompt: str) -> dict:
+    """One streaming /v1/responses request: the data of its response.completed event."""
+    url = f"http://{server.server_host}:{server.server_port}/v1/responses"
+    res = requests.post(url, json={
+        "model": "test", "input": prompt, "max_output_tokens": n_predict,
+        "temperature": 0.0, "stream": True,
+    }, stream=True, timeout=600)
+    assert res.status_code == 200, res.text
+    completed = None
+    for raw in res.iter_lines():
+        line = raw.decode("utf-8")
+        if line.startswith("data: "):
+            data = json.loads(line[6:])
+            if data.get("type") == "response.completed":
+                completed = data
+    assert completed is not None, "the stream never reached response.completed"
+    return completed
+
+
+def test_a_streamed_response_carries_the_preempt_record_where_a_plain_one_does():
+    # what a client keeps from a streamed /v1/responses is data["response"], so the record has to be in that object, the same place the non-streamed body carries it
+    os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8"
+    _start(n_ctx=512)
+
+    completed = _stream_responses(24, _PROMPT_A)
+
+    assert completed["response"]["preempt"]["parks"] >= 1, completed["response"]
+    assert completed["response"]["preempt"]["recomputes"] == 0, completed["response"]
+
+    plain = server.make_request("POST", "/v1/responses", data={
+        "model": "test", "input": _PROMPT_B, "max_output_tokens": 4, "temperature": 0.0,
+    })
+    assert plain.status_code == 200, plain.body
+    assert sorted(plain.body["preempt"]) == sorted(completed["response"]["preempt"]) == ["parks", "recomputes"]
+
+
 def test_two_image_chats_that_outgrow_the_parking_budget_both_finish():
     # a media chunk could not be parked by recompute, so with the host budget spent nothing could be parked at all and the pool overflowing ended both chats. The chunk comes back the way it went in: re-encoded off the task, its cells reserved whole
     os.environ["LLAMA_MEDIA_MARKER"] = "<__media__>"