From 90492a1fc536f9a00ca8c5962bbdc1b3bacc3faa Mon Sep 17 00:00:00 2001 From: Brock Tice Date: Mon, 12 Oct 2026 06:52:54 -0500 Subject: [PATCH] --batch: a request's "priority" orders the waiting (a free slot and the next prompt read go to the highest; a long read gives way to a higher one) The field SGLang and vLLM read (an integer, higher first; absent = 0, the order as before). No running request is paused. Requests waiting for a slot never hold back the control lines, so a read that gave way (BYIELD) can always go on and free its slot. Co-Authored-By: Claude Opus 5.5 --- docs/BATCHING.md | 5 +++ serve/server.py | 61 +++++++++++++++++++++-------- serve/test_parallel.py | 88 +++++++++++++++++++++++++++++++++++++++--- 3 files changed, 132 insertions(+), 22 deletions(-) diff --git a/docs/BATCHING.md b/docs/BATCHING.md index 69e5e2f7e28..91babc4ca43 100644 --- a/docs/BATCHING.md +++ b/docs/BATCHING.md @@ -133,6 +133,11 @@ about 10-25% speed per request on this card". `--parallel N` is honoured as aske half as long is waiting, the server sends `BYIELD`; at its next chunk boundary the long read stops, the part read so far is copied into a slot, the short request is admitted, and the long one then goes on from its slot with the same chunks (at most twice per request). +- **Priority**: a request's `"priority"` (an integer, higher first; the field SGLang and vLLM read, so a gateway such + as LiteLLM passes it on) orders the waiting. A free slot goes to the highest priority waiting for one, the control + lines (the next prompt read) to the highest priority waiting for them, and a long read gives way to a waiting + request of a higher priority whatever its length - but never to a much shorter one of a lower priority. A running + request is never paused for it. Without the field every request is 0, the order is as before. - **Each slot is a conversation cache.** A finished slot keeps what it holds (the prompt, the answer, and the checkpoint at the prompt's last turn boundary); the next turn of that conversation goes to that slot and the engine copies its state back (50-60 ms for a short conversation) instead of reading the history again - also for diff --git a/serve/server.py b/serve/server.py index 5ca7f4e4dab..29384cb261f 100644 --- a/serve/server.py +++ b/serve/server.py @@ -939,7 +939,9 @@ def __init__(self, exe: str, args: list[str], cwd: str | None = None, log: str | self.slot_cv = threading.Condition() self.waiting = 0 # requests waiting for the control lines (ctl) self.admitting = 0 # #1603: ... and those that gave the lines back to wait for a slot - self.wait_lens: list[list[int]] = [] # ... their prompt lengths (a long read gives way to short ones) + self.admit_prios: list[int] = [] # ... their priorities (a free slot goes to the highest) + self.wait_lens: list[list] = [] # ... [prompt length, priority, after_epoch] (a long read gives + # way to short ones; a higher priority goes first) self.ctl_epoch = 0 # how often the control lines were taken self.ctl = threading.Lock() # one admission or solo request on the control lines at a time self.gen = self.__dict__.get("gen", 0) + 1 # which engine process this is (a request notes its own) @@ -1416,10 +1418,11 @@ def _reap_orphan_slots(self) -> list[int]: YIELDS_MAX = 2 # #656: how often one request's prompt read gives way to a shorter waiting one - def _take_control(self, cancel, plen: int, after_epoch: int | None = None): + def _take_control(self, cancel, plen: int, after_epoch: int | None = None, prio: int = 0): """Waits for the control lines (one prompt read at a time), yielding None heartbeats; False when cancelled. - `after_epoch`: a request whose read gave way lets the requests waiting then go first.""" - entry = [plen] + `after_epoch`: a request whose read gave way lets the requests waiting then go first. `prio`: while a request + of a higher priority waits (and is not itself letting others go first), this one does not take the lines.""" + entry = [plen, prio, after_epoch] with self.slot_cv: self.waiting += 1 self.wait_lens.append(entry) @@ -1432,7 +1435,8 @@ def _take_control(self, cancel, plen: int, after_epoch: int | None = None): # restart is under way it keeps waiting and goes on with the new engine. raise EngineDied("the engine stopped while this request waited; it was not sent") with self.slot_cv: - turn = after_epoch is None or self.ctl_epoch > after_epoch or self.waiting <= 1 + turn = ((after_epoch is None or self.ctl_epoch > after_epoch or self.waiting <= 1) and + not self._outranked(prio)) if turn and self.ctl.acquire(timeout=0.5): if self.alive(): break @@ -1469,10 +1473,21 @@ def _may_go_solo(self, left: int, times: int, embeddings) -> bool: with self.slot_cv: return sum(1 for b in self.slot_busy if b) == 1 and self.waiting == 0 - def _shorter_waiting(self, plen: int) -> bool: - """#656: someone waits for the control lines with a prompt under half this one's: worth giving way to.""" + def _shorter_waiting(self, plen: int, prio: int = 0) -> bool: + """#656: someone of the same or a higher priority waits for the control lines with a prompt under half this + one's, or someone of a higher priority waits at all: worth giving way to.""" with self.slot_cv: - return any(e[0] * 2 <= plen for e in self.wait_lens) + return any(e[1] > prio or (e[1] == prio and e[0] * 2 <= plen) for e in self.wait_lens) + + def _outranked(self, prio: int) -> bool: + """A request of a higher priority waits for the control lines and is not letting others go first. The caller + holds slot_cv. Requests waiting for a slot do not count: a read that gave way (BYIELD) holds a slot and needs + the lines to go on and free it, whatever its priority.""" + return any(e[1] > prio and (e[2] is None or self.ctl_epoch > e[2]) for e in self.wait_lens) + + def _free_slot(self, prompt: list[int], prio: int) -> int | None: + """pick_slot, unless a request of a higher priority waits for a slot (the caller holds slot_cv).""" + return None if any(p > prio for p in self.admit_prios) else self.pick_slot(prompt) def generate_batched(self, ids, max_new, sampling, cancel, embeddings=None): """--batch. Alone (no slot busy, nobody waiting): the solo path (GEN, with drafts: the fastest stream) @@ -1488,12 +1503,13 @@ 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 {}) + prio = request_priority(sampling) out: list[int] = [] pending: list[int] = [] prompt, left = list(ids), int(max_new) reused0 = None # the first read's reused tokens: a continued leg's count includes this request's own output carry = {} # #1869: the figures of the legs that ended in a hand-back to the solo path (the final DONE adds to them) - ok = yield from self._take_control(cancel, len(prompt)) + ok = yield from self._take_control(cancel, len(prompt), prio=prio) if not ok: return holding, born = True, self.gen # (the engine it now has the control lines of) @@ -1506,7 +1522,7 @@ def generate_batched(self, ids, max_new, sampling, cancel, embeddings=None): try: while True: # a request in a slot that is left alone goes back to the solo path if not holding: - ok = yield from self._take_control(cancel, len(prompt)) + ok = yield from self._take_control(cancel, len(prompt), prio=prio) if not ok: return holding, born = True, self.gen # (the engine it now has the control lines of) @@ -1526,7 +1542,7 @@ def others(): def give_way(): # at a PP line, and between two far apart (the engine takes it at its next boundary) nonlocal reserved if (reserved is None and not out and not embeddings and yields < self.YIELDS_MAX and - self._shorter_waiting(len(prompt))): + self._shorter_waiting(len(prompt), prio)): with self.slot_cv: # the slot the part read will wait in reserved = self.pick_slot(prompt) if reserved is not None: @@ -1576,17 +1592,20 @@ def give_way(): # at a PP line, and between two far apart (the engine takes it # long, short at 0.25 s steps under parallel: 2 (ENGINE_REVIEW finding 1, reproduced in 0.1.41). while True: with self.slot_cv: - slot = self.pick_slot(prompt) + slot = self._free_slot(prompt, prio) if slot is not None: self.slot_busy[slot] = True self.slot_claims.add(slot) if admit_wait: admit_wait = False self.admitting = max(0, self.admitting - 1) + self.admit_prios.remove(prio) + self.slot_cv.notify_all() # the next priority down may take a free slot now break if not admit_wait: # #1603: counted while it waits for a slot (it holds no admit_wait, admit_since = True, time.monotonic() # control lines then) self.admitting += 1 + self.admit_prios.append(prio) if holding: self.ctl.release() holding = False @@ -1602,12 +1621,12 @@ def give_way(): # at a PP line, and between two far apart (the engine takes it f"for a free slot (busy: {busy}, in use: {sorted(self.slot_claims)}, " f"ending: {sorted(self.slot_releasing)}) (issue #1603)", flush=True) self._reap_orphan_slots() - if self.pick_slot(prompt) is None: + if self._free_slot(prompt, prio) is None: self.slot_cv.wait(timeout=1.0) if cancel.is_set(): return if not holding: - ok = yield from self._take_control(cancel, len(prompt)) + ok = yield from self._take_control(cancel, len(prompt), prio=prio) if not ok: return holding, born = True, self.gen @@ -1619,7 +1638,7 @@ def give_way(): # at a PP line, and between two far apart (the engine takes it holding = False with self.slot_cv: epoch = self.ctl_epoch - ok = yield from self._take_control(cancel, len(prompt), after_epoch=epoch) + ok = yield from self._take_control(cancel, len(prompt), after_epoch=epoch, prio=prio) if not ok: return holding, born = True, self.gen # (the engine it now has the control lines of) @@ -1640,7 +1659,7 @@ def give_way(): # at a PP line, and between two far apart (the engine takes it def ask(): # as give_way above, for the slot this admission reads into nonlocal asked if (not asked and not embeddings and yields < self.YIELDS_MAX and - self._shorter_waiting(len(prompt))): + self._shorter_waiting(len(prompt), prio)): self._send(f"BYIELD {slot}") asked = True try: @@ -1775,6 +1794,8 @@ def ask(): # as give_way above, for the slot this admission reads into if admit_wait: with self.slot_cv: self.admitting = max(0, self.admitting - 1) + self.admit_prios.remove(prio) + self.slot_cv.notify_all() # a lower priority may take the slot this one waited for if reserved is not None: with self.slot_cv: self.slot_busy[reserved] = False @@ -2922,6 +2943,14 @@ def push(self, t: int) -> str: return delta +def request_priority(req: dict | None) -> int: + """The request's `priority` (an integer, higher first, as SGLang's and vLLM's field; absent or not an integer: 0). + With --batch it orders the requests waiting for the control lines and for a slot; a running request is never + paused for it, and requests without it are served as before.""" + p = (req or {}).get("priority") + return p if isinstance(p, int) and not isinstance(p, bool) else 0 + + OPENAI_MAX_STOP = 4 # OpenAI's limit on `stop` diff --git a/serve/test_parallel.py b/serve/test_parallel.py index 1a638919eb1..01639cd6862 100644 --- a/serve/test_parallel.py +++ b/serve/test_parallel.py @@ -16,7 +16,7 @@ from unittest import mock from serve.frontend import ChatTemplate -from serve.server import ByteTokenizer, Service, StrataEngine, engine_args, parallel_args, serve +from serve.server import ByteTokenizer, Service, StrataEngine, engine_args, parallel_args, request_priority, serve # The fake engine: one token every STEP seconds per active slot (a "window" serves every active slot at once); # GEN (solo) streams T lines; BGEN reads the prompt (one T line, DONE), answers BADM and continues in the slot. @@ -222,6 +222,25 @@ def said(args, env): self.assertEqual(args[-2:], ["--batch", "2"]) +class Priority(unittest.TestCase): + def test_request_priority(self): + self.assertEqual(request_priority({"priority": 7}), 7) + self.assertEqual(request_priority({"priority": -3}), -3) + for body in ({}, None, {"priority": True}, {"priority": "7"}, {"priority": 1.5}, {"priority": None}): + self.assertEqual(request_priority(body), 0, body) + + def test_outranked_skips_a_request_letting_others_go_first(self): + e = StrataEngine("missing-executable", [], lazy=True) + e.ctl_epoch = 4 + e.wait_lens = [[100, 5, None]] + self.assertTrue(e._outranked(0)) + self.assertFalse(e._outranked(5)) + e.wait_lens = [[100, 5, 4]] # it gave way at epoch 4 and nobody has taken the lines since + self.assertFalse(e._outranked(0)) + e.ctl_epoch = 5 + self.assertTrue(e._outranked(0)) + + class PickSlot(unittest.TestCase): def engine(self, n): e = StrataEngine("missing-executable", [], lazy=True) @@ -307,8 +326,9 @@ def tearDown(self): if getattr(self, "tmp", None): self.tmp.cleanup() - def chat(self, text, max_tokens=64): - body = {"messages": [{"role": "user", "content": text}], "max_tokens": max_tokens, "reasoning_effort": "none"} + def chat(self, text, max_tokens=64, **extra): + body = {"messages": [{"role": "user", "content": text}], "max_tokens": max_tokens, "reasoning_effort": "none", + **extra} 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: @@ -469,13 +489,14 @@ def go(t): active = [int(x.split()[3]) for x in self.log.read_text().splitlines() if x.startswith("BGEN")] self.assertTrue(all(a <= 1 for a in active), active) - def race(self, first, second, delay=0.15, max_tokens=64): - """`first` sent, `second` after `delay` seconds; -> {text: (reply, seconds it took)}""" + def race(self, first, second, delay=0.15, max_tokens=64, priority=None): + """`first` sent, `second` after `delay` seconds; -> {text: (reply, seconds it took)}. `priority`: {text: it}""" res = {} def go(t): t0 = time.time() - r = self.chat(t, max_tokens=max_tokens) + extra = {"priority": priority[t]} if priority and t in priority else {} + r = self.chat(t, max_tokens=max_tokens, **extra) res[t] = (r["choices"][0]["message"]["content"], time.time() - t0) a = threading.Thread(target=go, args=(first,)) a.start() @@ -591,6 +612,61 @@ def test_no_way_given_to_a_long_one_between_far_apart_progress_lines(self): self.assertEqual(sorted(r[0] for r in res.values()), ["ok, done."] * 2) self.assertFalse(any(x.startswith("YIELD ") for x in self.log.read_text().splitlines())) + def test_a_long_read_gives_way_to_a_higher_priority(self): + """A higher priority is worth giving way to whatever its length; a lower one only when much shorter (#656).""" + self.start(2) + long_q, lung_q = "long " * 600, "lung " * 600 + res = self.race(long_q, lung_q, priority={lung_q: 5}) + self.assertEqual(sorted(r[0] for r in res.values()), ["ok, done."] * 2) + self.assertLess(res[lung_q][1], res[long_q][1] - 0.5, res) + self.assertTrue(any(x.startswith("YIELD ") for x in self.log.read_text().splitlines())) + + def test_no_way_given_to_a_lower_priority(self): + self.start(2) + long_q = "long " * 600 + res = self.race(long_q, "short", priority={long_q: 5}) + self.assertEqual(sorted(r[0] for r in res.values()), ["ok, done."] * 2) + self.assertFalse(any(x.startswith("YIELD ") for x in self.log.read_text().splitlines())) + + def test_a_free_slot_goes_to_the_highest_priority_waiting(self): + """Both slots decode; a request without a priority waits, then one with a higher priority: the first slot that + frees goes to the later, higher one. Every request is answered.""" + self.start(2) + texts = ["first LONGREPLY", "second LONGREPLY", "low", "high priority " * 10] + prios = {texts[3]: 5} + res, errors = {}, [] + + def go(t): + try: + extra = {"priority": prios[t]} if t in prios else {} + res[t] = self.chat(t, max_tokens=200, **extra)["choices"][0]["message"]["content"] + except Exception as e: # noqa: BLE001 - reported below + errors.append((t[:6], repr(e))) + threads = [] + for t in texts: + th = threading.Thread(target=go, args=(t,)) + th.start() + threads.append(th) + time.sleep(0.1) + for th in threads: + th.join(45) + self.assertEqual(errors, []) + self.assertEqual(len(res), 4) + # the fake logs each admission's prompt length: the high one (the longer prompt) was admitted before the low one + sizes = [int(x.split()[2]) for x in self.log.read_text().splitlines() if x.startswith("BGEN")] + self.assertGreater(sizes[-2], sizes[-1], sizes) + self.assertFalse(any(self.engine.slot_busy)) + self.assertEqual(self.engine.admit_prios, []) + + def test_no_way_given_to_a_lower_priority_between_far_apart_progress_lines(self): + """The ask between PP lines (#2170) keeps the priority rule: a much shorter request of a lower priority gets no + BYIELD there either.""" + self.start(2, more=["--pp-every", "80"]) + long_q = "long " * 600 + res = self.race(long_q, "short", priority={long_q: 5}) + self.assertEqual(sorted(r[0] for r in res.values()), ["ok, done."] * 2) + self.assertFalse(any(x.startswith("YIELD ") for x in self.log.read_text().splitlines())) + def test_left_alone_it_goes_back_to_the_solo_path(self): """A request decoding in a slot whose neighbour finished goes back to the solo path (MTP drafts), continued from its slot - only with an engine that keeps the slots' conversations (INFO slot_cache=1)."""