serve: logprobs with "parallel" (batch slots) for one-token decisions; 400 for longer replies

With "parallel" set, generate() goes to generate_batched(), whose control-line reader (_control) had no LP
handling: the engine wrote the LP line (BGEN is read as GEN 1, logprobs=K included) and the server dropped it,
so every reply came back with "logprobs": null.

- _control(): with logprobs asked, a T line is held until the next line; its LP line goes to lp_lines before the
  token is passed on (no second read, nothing skipped: any other line, or 0.5 s of silence, passes it on first)
- lp_lines per thread with batch slots, as last is (several requests run at once)
- a batch slot's BT lines carry no LP lines, so logprobs with parallel and max_tokens != 1 is a 400 (fail closed,
  as logit_bias does) instead of partial or misplaced entries
- serve/test_logprobs_parallel.py: the fake engine of test_parallel with LP lines; one decision, four at once
  (each gets its own), not asked, longer replies refused, an engine without LP lines

Contributed by an AI agent (Ren, Claude) on behalf of karls0r; measured on 2x RTX 3090, layer split.

Author: karls0r

serve/server.py

@@ -614,6 +614,24 @@ def last(self, value):
             tl = self.__dict__.setdefault("_tl", threading.local())
             tl.last = value
 
+    # `lp_lines`: the LP lines of the tokens of the request being read (logprobs), taken by LogprobEntries on the
+    # thread that reads the request's tokens.  With batch slots several requests run at once: one deque per thread.
+    @property
+    def lp_lines(self):
+        if self.batch:
+            tl = self.__dict__.setdefault("_tl", threading.local())
+            if not hasattr(tl, "lp_lines"):
+                tl.lp_lines = collections.deque()
+            return tl.lp_lines
+        return self.__dict__.setdefault("_lp_lines", collections.deque())
+
+    @lp_lines.setter
+    def lp_lines(self, value):
+        if self.batch:
+            self.__dict__.setdefault("_tl", threading.local()).lp_lines = value
+        else:
+            self.__dict__["_lp_lines"] = value
+
     def __init__(self, exe: str, args: list[str], cwd: str | None = None, log: str | None = None,
                  env: dict | None = None, lazy: bool = False, cpus: list[int] | None = None):
         self.spawn = (exe, list(args), cwd, log, env, False, cpus)   # to start it again after it died (issue #27)
@@ -1035,18 +1053,34 @@ def _send(self, text: str):
         except (OSError, AttributeError):                # the pipe is gone (or a restart's close() took the process):
             raise EngineDied(f"the engine stopped unexpectedly (exit code {self.exit_code()})") from None
 
-    def _control(self, cancel, on_token, stop_when=None):
+    def _control(self, cancel, on_token, stop_when=None, want_lp=False):
         """Reads the control lines of the request on them (GEN / BGEN), yielding None heartbeats.  Calls on_token(id)
         for each `T`; returns ("done", None) at DONE, or ("badm", continues) at BADM (after DONE).  `stop_when()`
-        true sends STOP once (the request is then read to its DONE)."""
+        true sends STOP once (the request is then read to its DONE).  `want_lp` (logprobs): a `T` is held until the
+        next line, and its `LP` line (the engine writes it right after the token, in the same flush) goes to lp_lines
+        before the token is passed on, so whoever takes the token finds it; any other line, or 0.5 s of silence,
+        passes the token on first (an engine without LP lines)."""
         stopped = False
         beat = time.monotonic()
+        held = None
+
+        def pass_on(t):
+            nonlocal stopped
+            on_token(t)
+            if not stopped and (cancel.is_set() or (stop_when is not None and stop_when())):
+                self._send("STOP")
+                stopped = True
         while True:
             try:
                 # A client can leave before the first PP/token.  Poll its Event as often as _take_control does;
                 # the ten-second HTTP heartbeat is independent of how promptly STOP reaches the engine.
                 line = self.lines.get(timeout=0.5)
             except queue.Empty:
+                if held is not None:
+                    t, held = held, None
+                    pass_on(t)
+                    yield False
+                    continue
                 if cancel.is_set() and not stopped:
                     self._send("STOP")
                     stopped = True
@@ -1059,11 +1093,19 @@ def _control(self, cancel, on_token, stop_when=None):
                 raise EngineDied(f"the engine stopped unexpectedly (exit code {self.exit_code()})")
             if not line.startswith(("PP ", "INFO")):
                 btrace("ctl<", self._ctl_mode, line.strip()[:60])
+            if held is not None:
+                t, held = held, None
+                if line.startswith("LP "):
+                    self.lp_lines.append((t, line))
+                pass_on(t)
+                yield False                              # a token is pending (not a heartbeat)
+                if line.startswith("LP "):
+                    continue
             if line.startswith("T "):
-                on_token(int(line[2:]))
-                if not stopped and (cancel.is_set() or (stop_when is not None and stop_when())):
-                    self._send("STOP")
-                    stopped = True
+                if want_lp:
+                    held = int(line[2:])
+                    continue
+                pass_on(int(line[2:]))
                 yield False                              # a token is pending (not a heartbeat)
             elif line.startswith("PP "):
                 f = line.split()
@@ -1255,6 +1297,8 @@ def generate_batched(self, ids, max_new, sampling, cancel, embeddings=None):
         self.progress, self.progress_ms, self.reused = None, 0, 0
         born = self.gen                                 # the engine process this request is sent to
         keys = self.sampling_keys(sampling or {})
+        want_lp = top_logprobs(sampling or {}) is not None
+        self.lp_lines = collections.deque()             # this request's (per thread with batch slots)
         out: list[int] = []
         pending: list[int] = []
         prompt, left = list(ids), int(max_new)
@@ -1287,7 +1331,7 @@ def generate_batched(self, ids, max_new, sampling, cancel, embeddings=None):
                     def others():
                         with self.slot_cv:
                             return self.waiting > 0
-                    for x in self._control(cancel, pending.append, stop_when=others):
+                    for x in self._control(cancel, pending.append, stop_when=others, want_lp=want_lp):
                         while pending:
                             t = pending.pop(0)
                             out.append(t)
@@ -1392,7 +1436,7 @@ def others():
                     self._ctl_mode, self._ctl_result, self._yielded = "batch", None, None
                     self._ctl_erred = False
                     asked = False
-                    for x in self._control(cancel, pending.append):
+                    for x in self._control(cancel, pending.append, want_lp=want_lp):
                         while pending:
                             t = pending.pop(0)
                             out.append(t)
@@ -5370,6 +5414,11 @@ def _openai(self, req):
             svc.load()
             max_req = max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 0)   # 0/-1: the rest
             bias_api.validate_request(req, svc.engine, len(getattr(svc.tok, "tokens", ())) or None)
+            if top_logprobs(req) is not None and getattr(svc.engine, "batch", 0) and max_req != 1:
+                # a batch slot's BT lines carry no LP lines: past its first token a request would get partial or
+                # misplaced entries.  One token (a typed decision) is read on the control lines and has its LP line.
+                raise ValueError('logprobs: with "parallel" (batch slots) only for max_tokens 1; '
+                                 'run without "parallel" for longer replies')
             use_mcp = req.get("strata_mcp") is True and svc.mcp is not None      # the web app's opt-in (serve/mcp.py)
             own = {t.get("name") for t in tools or [] if isinstance(t, dict)}   # #592: a second line of defence
             if use_mcp:

serve/test_logprobs_parallel.py

@@ -0,0 +1,101 @@
+"""Logprobs with "parallel" (batch slots): python -m unittest serve.test_logprobs_parallel.
+
+The fake engine of serve/test_parallel.py behind the real StrataEngine and Service, writing an LP line after each
+T line when the GEN / BGEN line asks for logprobs.  A typed decision (max_tokens 1) is read on the control lines (an
+admission: BGEN, one T, DONE, BADM 0) and gets its logprobs; longer replies are refused with a 400, since the slot's
+BT lines carry none."""
+import json
+import threading
+import unittest
+import urllib.error
+import urllib.request
+from unittest import mock
+
+import serve.test_parallel as tp
+
+# LP after each T when asked: the token's log-probability is -(prompt length) / 1000 (so two requests can be told
+# apart) and "A" (65) is its one alternative at -1.5
+_T_LINE = '            print(f"T {t}", flush=True)\n'
+assert tp.FAKE_BATCH.count(_T_LINE) == 1
+FAKE_LP = tp.FAKE_BATCH.replace(_T_LINE, _T_LINE + (
+    '            if " logprobs=" in line:\n'
+    '                print(f"LP {-len(ids) / 1000:.6f} 65:-1.500000", flush=True)\n'))
+
+
+class LogprobsWithBatchSlots(unittest.TestCase):
+    start, tearDown, get = tp.ParallelService.start, tp.ParallelService.tearDown, tp.ParallelService.get
+
+    def chat(self, text, max_tokens=1, logprobs=True):
+        body = {"messages": [{"role": "user", "content": text}], "max_tokens": max_tokens,
+                "reasoning_effort": "none"}
+        if logprobs:
+            body |= {"logprobs": True, "top_logprobs": 1}
+        req = urllib.request.Request(self.base + "/v1/chat/completions", data=json.dumps(body).encode(),
+                                     headers={"Content-Type": "application/json"})
+        with urllib.request.urlopen(req, timeout=60) as r:
+            return json.loads(r.read().decode())
+
+    def start_lp(self, fake=FAKE_LP):
+        with mock.patch.object(tp, "FAKE_BATCH", fake):
+            self.start(2)
+
+    @staticmethod
+    def entries(reply):
+        return ((reply["choices"][0].get("logprobs") or {}).get("content")) or []
+
+    def test_a_decision_gets_its_logprobs(self):
+        self.start_lp()
+        r = self.chat("Is it A or B?")
+        self.assertEqual(r["choices"][0]["message"]["content"], "o")
+        e = self.entries(r)
+        self.assertEqual(len(e), 1)
+        self.assertEqual(e[0]["token"], "o")
+        self.assertLess(e[0]["logprob"], 0)
+        self.assertEqual(e[0]["top_logprobs"], [{"token": "A", "logprob": -1.5, "bytes": [65]}])
+        self.assertTrue(any(x.startswith("BGEN") for x in self.log.read_text().splitlines()))   # an admission
+
+    def test_two_at_once_each_get_their_own(self):
+        self.start_lp()
+        short, long_ = "A or B?", "A or B? " + "x" * 300
+        alone = {t: self.entries(self.chat(t))[0]["logprob"] for t in (short, long_)}
+        self.assertNotEqual(alone[short], alone[long_])
+        for _ in range(5):
+            got, errs = {}, []
+
+            def run(t):
+                try:
+                    got[t] = self.entries(self.chat(t))[0]["logprob"]
+                except Exception as exc:          # noqa: BLE001 - reported below
+                    errs.append(repr(exc))
+            th = [threading.Thread(target=run, args=(t,)) for t in (short, long_, short, long_)]
+            for x in th:
+                x.start()
+            for x in th:
+                x.join(60)
+            self.assertEqual(errs, [])
+            self.assertEqual(got, alone)
+
+    def test_not_asked_nothing_changes(self):
+        self.start_lp()
+        r = self.chat("Is it A or B?", logprobs=False)
+        self.assertEqual(r["choices"][0]["message"]["content"], "o")
+        self.assertIsNone(r["choices"][0].get("logprobs"))
+
+    def test_longer_replies_are_refused_with_batch_slots(self):
+        self.start_lp()
+        with self.assertRaises(urllib.error.HTTPError) as cm:
+            self.chat("Is it A or B?", max_tokens=8)
+        self.assertEqual(cm.exception.code, 400)
+        self.assertIn("max_tokens 1", cm.exception.read().decode())
+        self.assertEqual(self.chat("still serving?", max_tokens=16, logprobs=False)["choices"][0]["message"]["content"],
+                         "ok, done.")
+
+    def test_an_engine_without_lp_lines(self):
+        self.start_lp(fake=tp.FAKE_BATCH)
+        r = self.chat("Is it A or B?")
+        self.assertEqual(r["choices"][0]["message"]["content"], "o")
+        self.assertEqual(self.entries(r), [])
+
+
+if __name__ == "__main__":
+    unittest.main()