Skip to content

--batch: request priority (a free slot and the next prompt read go to the highest priority waiting) #2109

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
brocktice wants to merge 1 commit into from
Closed

--batch: request priority (a free slot and the next prompt read go to the highest priority waiting) #2109

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions docs/BATCHING.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
61 changes: 45 additions & 16 deletions serve/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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`


Expand Down
88 changes: 82 additions & 6 deletions serve/test_parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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)."""
Expand Down