perf: batch decode stream callbacks into one dispatch per step
- add BatchedStreamCallback sink type: TaskManager resolves a decode step's (task_id, token) events under one lock and delivers each sink a single list instead of one call per token - keep the plain Callable[[str]] callback contract: per-token callbacks still receive one call per event, and invoke_callback/cancel_task wrap single events for batched sinks - collect aborted, text, and finish STOP events in the scheduler decode loop and dispatch once per step instead of once per token - register one _ResultSink per generate call (replacing per-task closures) so GenerateResult takes its lock and wakes waiters once per step, with late-bind replay for tasks that start decoding before add_task returns their id - apply GenerateResult batches under a single condition hold via append_batch; append delegates to it - update engine test fakes to the batched contract and add coverage for event grouping, single-event dispatch, cancel STOP, and late-bind replay Benchmark: NVIDIA L20 (idle), CUDA 12.8, torch 2.11.0+cu128, 1.2B bf16 checkpoint, prompt 512, 256 greedy tokens, CUDA graph on, serving-level decode, 3 trials - batch 32: 7.808 -> 7.506 ms/token (4098 -> 4263 batch tok/s, +4.0%) - batch 1/8: unchanged within noise (3.768 -> 3.797 / 4.699 -> 4.607 ms/token) - full suite: 896 passed
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
"""Unit tests for GenerateResult accumulator and InferenceEngine.generate()."""
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
import threading
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -8,7 +9,12 @@ import pytest
|
||||
|
||||
from astrai.extension import TorchNativeBackend, attn_backend
|
||||
from astrai.inference import STOP
|
||||
from astrai.inference.engine import GenerateResult, InferenceEngine, build_engine
|
||||
from astrai.inference.engine import (
|
||||
GenerateResult,
|
||||
InferenceEngine,
|
||||
_ResultSink,
|
||||
build_engine,
|
||||
)
|
||||
from tests.helpers import FakeTokenizer, make_model
|
||||
|
||||
|
||||
@@ -50,6 +56,33 @@ def test_result_stop_does_not_double_count():
|
||||
assert r._completed == 1
|
||||
|
||||
|
||||
def test_result_append_batch_updates_state_in_one_commit():
|
||||
r = GenerateResult(count=2)
|
||||
r.append_batch([(0, "he"), (1, "wo"), (0, "llo"), (1, "rld")])
|
||||
r.append_batch([(0, STOP), (1, STOP)])
|
||||
assert r.results == ["hello", "world"]
|
||||
assert r._completed == 2
|
||||
assert r.pop_all() == [
|
||||
(0, "he"),
|
||||
(1, "wo"),
|
||||
(0, "llo"),
|
||||
(1, "rld"),
|
||||
(0, STOP),
|
||||
(1, STOP),
|
||||
]
|
||||
|
||||
|
||||
def test_result_sink_replays_events_arriving_before_bind():
|
||||
r = GenerateResult(count=1)
|
||||
sink = _ResultSink(r)
|
||||
sink([("t0", "he")]) # task id not bound yet: buffered, not applied
|
||||
assert r.results == [""]
|
||||
sink.bind("t0", 0)
|
||||
sink([("t0", "llo"), ("t0", STOP)])
|
||||
assert r.results == ["hello"]
|
||||
assert r._completed == 1
|
||||
|
||||
|
||||
def test_result_pop_all_returns_and_clears():
|
||||
r = GenerateResult(count=2)
|
||||
r.append("a", 0)
|
||||
@@ -118,8 +151,8 @@ def test_engine_generate_non_streaming_single():
|
||||
|
||||
def fake_add(prompt, **kw):
|
||||
cb = kw["stream_callback"]
|
||||
cb("response")
|
||||
cb(STOP)
|
||||
cb([("task-1", "response"), ("task-1", STOP)])
|
||||
return "task-1"
|
||||
|
||||
instance.add_task.side_effect = fake_add
|
||||
instance.remove_task.return_value = []
|
||||
@@ -136,6 +169,7 @@ def test_engine_generate_streaming_yields_tokens():
|
||||
|
||||
def capture_cb(prompt, **kw):
|
||||
callbacks_saved.append(kw.get("stream_callback"))
|
||||
return "task-0"
|
||||
|
||||
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
|
||||
instance = MockSched.return_value
|
||||
@@ -146,9 +180,9 @@ def test_engine_generate_streaming_yields_tokens():
|
||||
gen = eng.generate("hello", stream=True)
|
||||
|
||||
cb = callbacks_saved[0]
|
||||
cb("t1")
|
||||
cb("t2")
|
||||
cb(STOP)
|
||||
cb([("task-0", "t1")])
|
||||
cb([("task-0", "t2")])
|
||||
cb([("task-0", STOP)])
|
||||
|
||||
tokens = list(gen)
|
||||
assert tokens == ["t1", "t2"]
|
||||
@@ -169,7 +203,7 @@ def test_engine_stream_close_cancels_unfinished_task():
|
||||
engine = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1)
|
||||
stream = engine.generate("hello", stream=True)
|
||||
|
||||
callbacks_saved[0]("t1")
|
||||
callbacks_saved[0]([("task-1", "t1")])
|
||||
assert next(stream) == "t1"
|
||||
stream.close()
|
||||
|
||||
@@ -182,6 +216,7 @@ def test_engine_generate_async_yields_tokens_until_stop():
|
||||
|
||||
def capture_cb(prompt, **kw):
|
||||
callbacks_saved.append(kw.get("stream_callback"))
|
||||
return "task-0"
|
||||
|
||||
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
|
||||
instance = MockSched.return_value
|
||||
@@ -198,9 +233,9 @@ def test_engine_generate_async_yields_tokens_until_stop():
|
||||
return out
|
||||
|
||||
cb = callbacks_saved[0]
|
||||
cb("t1")
|
||||
cb("t2")
|
||||
cb(STOP)
|
||||
cb([("task-0", "t1")])
|
||||
cb([("task-0", "t2")])
|
||||
cb([("task-0", STOP)])
|
||||
|
||||
assert asyncio.run(collect()) == ["t1", "t2"]
|
||||
|
||||
@@ -219,7 +254,7 @@ def test_engine_async_close_cancels_unfinished_task():
|
||||
|
||||
engine = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1)
|
||||
stream = engine.generate_async("hello")
|
||||
callbacks_saved[0]("t1")
|
||||
callbacks_saved[0]([("task-1", "t1")])
|
||||
|
||||
async def consume_then_close():
|
||||
assert await anext(stream) == "t1"
|
||||
@@ -233,20 +268,25 @@ def test_engine_async_close_cancels_unfinished_task():
|
||||
def test_engine_generate_non_streaming_batch():
|
||||
mock_model, mock_tokenizer = _make_engine_mocks(decode="r")
|
||||
|
||||
counter = itertools.count()
|
||||
task_ids = []
|
||||
|
||||
def fake_add(prompt, **kw):
|
||||
cb = kw["stream_callback"]
|
||||
tid = f"task-{next(counter)}"
|
||||
task_ids.append(tid)
|
||||
cb([(tid, "r"), (tid, STOP)])
|
||||
return tid
|
||||
|
||||
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
|
||||
instance = MockSched.return_value
|
||||
|
||||
def fake_add(prompt, **kw):
|
||||
cb = kw["stream_callback"]
|
||||
cb("r")
|
||||
cb(STOP)
|
||||
|
||||
instance.add_task.side_effect = fake_add
|
||||
instance.remove_task.return_value = []
|
||||
|
||||
eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=2)
|
||||
results = eng.generate(["hello", "world"])
|
||||
assert results == ["r", "r"]
|
||||
assert task_ids == ["task-0", "task-1"]
|
||||
|
||||
|
||||
def test_engine_generate_zero_max_tokens_returns_empty():
|
||||
@@ -294,7 +334,7 @@ def test_generate_captures_calling_backend_context():
|
||||
|
||||
def fake_add(prompt, **kwargs):
|
||||
captured.append(kwargs["backend"])
|
||||
kwargs["stream_callback"](STOP)
|
||||
kwargs["stream_callback"]([("task", STOP)])
|
||||
return "task"
|
||||
|
||||
instance.add_task.side_effect = fake_add
|
||||
@@ -328,9 +368,12 @@ def test_build_engine_passes_engine_kwargs_through():
|
||||
model, _ = make_model("cpu", max_position_embeddings=64)
|
||||
backend = TorchNativeBackend()
|
||||
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
|
||||
MockSched.return_value.add_task.side_effect = lambda *args, **k: (
|
||||
k["stream_callback"](STOP) or "task"
|
||||
)
|
||||
|
||||
def fake_add(*args, **k):
|
||||
k["stream_callback"]([("task", STOP)])
|
||||
return "task"
|
||||
|
||||
MockSched.return_value.add_task.side_effect = fake_add
|
||||
engine = build_engine(
|
||||
model=model,
|
||||
tokenizer=FakeTokenizer(),
|
||||
|
||||
@@ -4,7 +4,23 @@ from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from astrai.inference import STOP, Task, TaskManager, TaskStatus
|
||||
from astrai.inference import (
|
||||
STOP,
|
||||
BatchedStreamCallback,
|
||||
Task,
|
||||
TaskManager,
|
||||
TaskStatus,
|
||||
)
|
||||
|
||||
|
||||
class RecordingSink(BatchedStreamCallback):
|
||||
"""Batch-aware callback capturing every dispatch as one batch."""
|
||||
|
||||
def __init__(self):
|
||||
self.batches = []
|
||||
|
||||
def __call__(self, events):
|
||||
self.batches.append(events)
|
||||
|
||||
|
||||
def _make_mock_tokenizer():
|
||||
@@ -217,3 +233,46 @@ def test_task_manager_cancel_active_task_delivers_stop_callback():
|
||||
immediate, cancelled = tm.cancel_task(task_id)
|
||||
assert cancelled and immediate == []
|
||||
assert received == [STOP]
|
||||
|
||||
|
||||
def test_invoke_callbacks_batches_sink_events_and_keeps_plain_per_token():
|
||||
tm = TaskManager(tokenizer=_make_mock_tokenizer())
|
||||
plain = []
|
||||
tid_plain = tm.add_task("plain", stream_callback=plain.append)
|
||||
sink = RecordingSink()
|
||||
tid_a = tm.add_task("sink a", stream_callback=sink)
|
||||
tid_b = tm.add_task("sink b", stream_callback=sink)
|
||||
|
||||
tm.invoke_callbacks(
|
||||
[
|
||||
(tid_a, "x"),
|
||||
(tid_plain, "p"),
|
||||
(tid_b, "y"),
|
||||
("unknown-task", "dropped"),
|
||||
(tid_a, STOP),
|
||||
]
|
||||
)
|
||||
|
||||
assert plain == ["p"]
|
||||
assert sink.batches == [[(tid_a, "x"), (tid_b, "y"), (tid_a, STOP)]]
|
||||
|
||||
|
||||
def test_invoke_callback_delivers_single_event_to_batched_sink():
|
||||
tm = TaskManager(tokenizer=_make_mock_tokenizer())
|
||||
sink = RecordingSink()
|
||||
task_id = tm.add_task("test", stream_callback=sink)
|
||||
|
||||
tm.invoke_callback(task_id, STOP)
|
||||
|
||||
assert sink.batches == [[(task_id, STOP)]]
|
||||
|
||||
|
||||
def test_cancel_delivers_batched_stop_to_sink():
|
||||
tm = TaskManager(tokenizer=_make_mock_tokenizer())
|
||||
sink = RecordingSink()
|
||||
task_id = tm.add_task("test", stream_callback=sink)
|
||||
|
||||
immediate, cancelled = tm.cancel_task(task_id)
|
||||
|
||||
assert cancelled
|
||||
assert sink.batches == [[(task_id, STOP)]]
|
||||
|
||||
Reference in New Issue
Block a user