diff --git a/packages/sie_server/src/sie_server/config/engine.py b/packages/sie_server/src/sie_server/config/engine.py index ac3aede46..0acab55d6 100644 --- a/packages/sie_server/src/sie_server/config/engine.py +++ b/packages/sie_server/src/sie_server/config/engine.py @@ -303,13 +303,16 @@ class EngineConfig(BaseSettings): Field( ge=0.0, description=( - "Accumulation window for the FIRST batch after a worker was " - "idle (#2874). An idle worker used to dispatch immediately " - "with whatever was pending, shredding bursty arrivals into a " - "train of small serialized forwards. With this window the " - "batcher coalesces until arrivals have been quiet for this " - "long (bounded by max_batch_wait_ms), so a lone request waits " - "at most this window. 0 restores immediate idle dispatch." + "Tail window for a burst that reaches an idle worker " + "(#2874, #373). One request pending on a batcher is " + "dispatched immediately, however many items it carries, and " + "is not held for this long. A follow-up waits only while " + "another request on that same batcher is still pending. A " + "finished request does not hold the next one, and another " + "LoRA's arrival does not hold this batcher. Several requests " + "already pending still form one batch. Bounded by " + "max_batch_wait_ms. 0 dispatches every idle arrival " + "immediately." ), ), ] = 3.0 diff --git a/packages/sie_server/src/sie_server/core/batcher.py b/packages/sie_server/src/sie_server/core/batcher.py index c4a5cf4ca..f49c1b2f9 100644 --- a/packages/sie_server/src/sie_server/core/batcher.py +++ b/packages/sie_server/src/sie_server/core/batcher.py @@ -237,9 +237,22 @@ def config(self) -> BatchConfig: @property def pending_count(self) -> int: - """Return number of pending requests.""" + """Return number of pending items. + + One submitted request occupies one entry per item. + """ return len(self._pending) + @property + def pending_request_count(self) -> int: + """Return number of distinct pending requests. + + ``submit_many`` appends one entry per item and reuses one metadata + object, so ``pending_count`` counts items. Callers that mean "one + request" (idle dispatch, #373) use this. + """ + return len({id(request.metadata) for request in self._pending}) + @property def pending_cost(self) -> int: """Return total cost in pending requests.""" diff --git a/packages/sie_server/src/sie_server/core/worker/model_worker.py b/packages/sie_server/src/sie_server/core/worker/model_worker.py index f33464355..e922edb3e 100644 --- a/packages/sie_server/src/sie_server/core/worker/model_worker.py +++ b/packages/sie_server/src/sie_server/core/worker/model_worker.py @@ -178,6 +178,27 @@ def _dispatch_width(adapter: object) -> int: return max(1, width) +async def _yield_ready_callbacks() -> None: + """Let ready producer callbacks schedule submissions before the idle recheck. + + Completing a Future queues this task behind work those callbacks create. + This is one scheduling checkpoint, not a drain of all incoming work: later + network arrivals or unfinished preparation may still miss the batch. + """ + loop = asyncio.get_running_loop() + ready: asyncio.Future[None] = loop.create_future() + + def wake() -> None: + if not ready.done(): + ready.set_result(None) + + handle = loop.call_soon(wake) + try: + await ready + finally: + handle.cancel() + + class ModelWorker: """Worker that batches and processes inference requests for a single model. @@ -348,6 +369,7 @@ def __init__( self._running = False self._stopping = False # True when graceful stop has begun self._process_task: asyncio.Task[None] | None = None + # ``stop()`` sets this too so an idle ``_get_next_batch_fcfs`` unblocks. self._work_available = asyncio.Event() self._stats = WorkerStats() @@ -1328,20 +1350,25 @@ async def _get_next_batch_fcfs( Selects the batcher whose first pending request has waited the longest. This ensures fairness across LoRAs - no LoRA starves even with low traffic. - When the worker was idle (had to wait for requests), a short - accumulation window (``idle_coalesce_ms``, #2874) caps the coalesce - wait instead of dispatching immediately: bursty arrivals at an idle - worker fuse into one batch rather than degenerating into a train of - small serialized forwards, while a lone request only waits the small - cap once arrivals stop. Setting ``idle_coalesce_ms=0`` restores the - legacy immediate dispatch. + When the worker was idle, one pending request is dispatched immediately + (#373), however many items it carries. It is not held for + ``idle_coalesce_ms``. A lone head yields one callback generation so + ready producers and the submissions they schedule can join it, without + waiting for a timer. Later arrivals may still miss this batch. A + follow-up is a tail only while another request + on that same batcher is still pending (arrived and not yet dispatched); + the idle window can then catch a straggler. A finished request is no + longer pending, so it does not hold the next one, and another LoRA's + arrival does not hold this batcher. Several requests already pending + when a batch is selected form one batch. ``idle_coalesce_ms=0`` + disables the tail window and dispatches every idle arrival immediately. Args: was_idle: Whether the worker was idle before this call. When True, - dispatches after the capped idle accumulation window (or - immediately when the window is 0). When False, uses the - normal timeout/coalesce mechanism to accumulate a proper - batch. + one pending request dispatches immediately and two or more on + the selected batcher use the capped idle window. When False, + uses the normal timeout/coalesce mechanism to accumulate a + proper batch. Returns: Tuple of (lora_name, batch, was_idle) where lora_name is None for @@ -1349,6 +1376,9 @@ async def _get_next_batch_fcfs( (was truly idle). """ while True: + if not self._running: + return None, FormattedBatch(items=[], metadata=[], total_cost=0), was_idle + oldest_lora: str | None = None oldest_time: float = float("inf") @@ -1363,14 +1393,25 @@ async def _get_next_batch_fcfs( if oldest_lora is not None or (oldest_lora is None and self._batchers[None].pending_count > 0): # Found a batcher with pending items - get batch from it selected_lora = oldest_lora if oldest_lora is not None else None + selected = self._batchers[selected_lora] idle_window_ms = self._config.idle_coalesce_ms - if was_idle and idle_window_ms > 0: - batch = await self._batchers[selected_lora].get_batch(coalesce_cap_ms=idle_window_ms) + if was_idle and idle_window_ms > 0 and selected.pending_request_count == 1: + # A ready transport callback may create the submission task. + # Give that task one opportunity before the lone-head check. + await _yield_ready_callbacks() + if not self._running: + return None, FormattedBatch(items=[], metadata=[], total_cost=0), was_idle + # Alone is one request, not one item: ``submit_many`` queues + # an entry per item. The distinct-request count is only the + # idle tail check. A busy worker, and ``idle_coalesce_ms=0``, + # dispatch without consulting it. + if was_idle and idle_window_ms > 0 and selected.pending_request_count > 1: + batch = await selected.get_batch(coalesce_cap_ms=idle_window_ms) else: - batch = await self._batchers[selected_lora].get_batch(immediate=was_idle) + batch = await selected.get_batch(immediate=was_idle) return selected_lora, batch, was_idle - # No batchers have pending items - worker is idle + # No batchers have pending items - worker is idle until submit(). was_idle = True # Admission latches the signal even if it precedes this wait. # Re-scan every batcher after clearing; queue state owns the work. @@ -1435,10 +1476,10 @@ async def _process_loop(self) -> None: await self._process_loop_concurrent() return - # Track idle state across iterations. When idle, the next batch is - # dispatched immediately (low-concurrency optimization). When busy - # (just finished inference + drain), we let BatchFormer's - # timeout/coalesce mechanism accumulate a proper batch. + # Track idle state across iterations. An idle worker dispatches one + # pending request immediately (#373). Two or more still pending on + # the selected batcher use the idle window. After a busy batch, + # BatchFormer's timeout/coalesce mechanism accumulates the next batch. was_idle = True while self._running: diff --git a/packages/sie_server/src/sie_server/core/worker/types.py b/packages/sie_server/src/sie_server/core/worker/types.py index 27d076e84..c3c07b137 100644 --- a/packages/sie_server/src/sie_server/core/worker/types.py +++ b/packages/sie_server/src/sie_server/core/worker/types.py @@ -216,15 +216,13 @@ class WorkerConfig: # instead of being shredded into several half-full ones. coalesce_ms: float = 15.0 coalesce_ratio: float = 0.5 - # Idle-dispatch accumulation window (#2874). An idle worker used to - # dispatch IMMEDIATELY with whatever was pending, so a burst arriving at - # an idle worker degenerated into a train of small serialized forwards - # (many batch-of-1 dispatches instead of one fused batch). With this - # window, the first batch after idleness coalesces arrivals until the - # queue has been quiet for ``idle_coalesce_ms`` (or the batch fills, or - # ``max_batch_wait_ms`` elapses since the first request). A lone request - # therefore waits at most this window — single-digit ms — before - # dispatch, and ``0`` restores the legacy immediate dispatch. + # Idle-dispatch tail window (#2874, #373). One request pending on an + # idle batcher is dispatched immediately, however many items it carries, + # and is not held for this long. A follow-up waits only while another + # request on that same batcher is still pending. A finished request does + # not hold the next one, and another LoRA's arrival does not hold this + # batcher. Several requests already pending when a batch is selected + # still form one batch. ``0`` dispatches every idle arrival immediately. idle_coalesce_ms: float = 3.0 max_queue_size: int = 1000 # Maximum pending items in queue (0 = unlimited) instrumentation: bool = False diff --git a/packages/sie_server/tests/core/test_worker_core.py b/packages/sie_server/tests/core/test_worker_core.py index d85682e22..3d5fc25d9 100644 --- a/packages/sie_server/tests/core/test_worker_core.py +++ b/packages/sie_server/tests/core/test_worker_core.py @@ -16,6 +16,54 @@ from sie_server.types.inputs import Item +@pytest.fixture(params=["asyncio", "uvloop"]) +def ready_loop_factory(request: pytest.FixtureRequest): + if request.param == "uvloop": + return pytest.importorskip("uvloop").new_event_loop + return asyncio.new_event_loop + + +@pytest.mark.parametrize("stage", ["before_start", "before_resolve", "after_resolve_before_resume"]) +def test_ready_callback_checkpoint_cancellation( + stage: str, monkeypatch: pytest.MonkeyPatch, ready_loop_factory +) -> None: + async def check() -> None: + """Cancellation cannot leave an owned callback or resolve a cancelled Future.""" + loop = asyncio.get_running_loop() + call_soon = loop.call_soon + handles: list[asyncio.Handle] = [] + errors: list[dict] = [] + old_handler = loop.get_exception_handler() + + def record_callback(callback, *args, context=None): + handle = call_soon(callback, *args, context=context) + if getattr(callback, "__module__", None) == model_worker_module.__name__: + handles.append(handle) + return handle + + monkeypatch.setattr(loop, "call_soon", record_callback) + loop.set_exception_handler(lambda _loop, context: errors.append(context)) + try: + task = asyncio.create_task(model_worker_module._yield_ready_callbacks()) + if stage == "before_start": + task.cancel() + elif stage == "before_resolve": + loop.call_soon(task.cancel) + else: + loop.call_soon(lambda: loop.call_soon(task.cancel)) + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0) + assert len(handles) == (0 if stage == "before_start" else 1) + assert all(handle.cancelled() for handle in handles) + assert errors == [] + finally: + loop.set_exception_handler(old_handler) + + with asyncio.Runner(loop_factory=ready_loop_factory) as runner: + runner.run(check()) + + class TestRequestMetadata: """Tests for RequestMetadata dataclass.""" @@ -76,8 +124,8 @@ def test_defaults(self) -> None: assert config.max_batch_wait_ms == 15 assert config.coalesce_ms == 15.0 assert config.coalesce_ratio == 0.5 - # Idle accumulation window (#2874): single-digit ms by default so an - # idle worker fuses bursts without materially delaying lone requests. + # Tail window for a staggered burst (#2874, #373). Default stays + # non-zero; a lone request is not charged it. assert config.idle_coalesce_ms == 3.0 def test_custom_values(self) -> None: @@ -586,11 +634,12 @@ async def test_idle_dispatch_no_wait(self, mock_adapter: MagicMock, tokenized_it @pytest.mark.asyncio async def test_staggered_burst_fuses_with_idle_window(self, mock_adapter: MagicMock) -> None: - """#2874: a burst arriving at an IDLE worker fuses into one batch. + """#373: the head runs alone, and a still-pending tail fuses behind it. - Before the idle accumulation window, an idle worker dispatched - immediately with whatever was pending, so staggered arrivals became a - train of serialized batch-of-1 forwards (the was_idle degeneracy). + The head is not charged ``idle_coalesce_ms``. A finished request does + not hold the next one. Two requests queued on the same batcher are a + tail, so the idle window holds them and a third arrival inside the + window joins that forward: two forwards, not one per request. """ call_sizes: list[int] = [] @@ -607,13 +656,205 @@ def counting_encode(items, output_types, **kwargs): max_batch_tokens=1000, max_batch_requests=64, max_batch_wait_ms=500, - coalesce_ms=200, - coalesce_ratio=0.5, - idle_coalesce_ms=60, + coalesce_ms=400, + coalesce_ratio=1.0, + idle_coalesce_ms=200, ) worker = ModelWorker(mock_adapter, config) await worker.start() + try: + started = time.monotonic() + first = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="text 0")], + ["dense"], + ) + await asyncio.wait_for(first, timeout=2.0) + head_ms = (time.monotonic() - started) * 1000 + # Must not have slept the 200ms window. 50ms leaves room for CI noise. + assert head_ms < 50, f"lone head waited {head_ms:.1f}ms" + assert call_sizes == [1], f"expected the head to run alone, got: {call_sizes}" + + # The head has finished, so it must not hold these. Queue two + # before the loop selects, then one more inside the 200ms window. + second = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="text 1")], + ["dense"], + ) + third = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="text 2")], + ["dense"], + ) + assert worker.pending_count == 2 + await asyncio.sleep(0.01) + assert worker.pending_count == 2, "tail was dispatched instead of held" + fourth = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="text 3")], + ["dense"], + ) + + await asyncio.gather(second, third, fourth) + + assert sum(call_sizes) == 4 + assert call_sizes == [1, 3], f"expected the tail to fuse, got: {call_sizes}" + + finally: + await worker.stop() + + @pytest.mark.asyncio + @pytest.mark.parametrize("warmup", [False, True]) + async def test_ready_submissions_join_the_idle_head(self, mock_adapter: MagicMock, warmup: bool) -> None: + """A runnable submit joins the head without adding a fixed idle delay.""" + call_sizes: list[int] = [] + + def counting_encode(items, output_types, **kwargs): + call_sizes.append(len(items)) + return EncodeOutput( + dense=np.array([[0.1, 0.2, 0.3]] * len(items)), + batch_size=len(items), + ) + + mock_adapter.encode.side_effect = counting_encode + worker = ModelWorker( + mock_adapter, + WorkerConfig( + max_batch_tokens=1000, + max_batch_requests=64, + max_batch_wait_ms=500, + coalesce_ms=400, + coalesce_ratio=1.0, + idle_coalesce_ms=3, + ), + ) + await worker.start() + + async def submit_ready(): + return await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="second")], + ["dense"], + ) + + try: + if warmup: + warmed = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="warmup")], + ["dense"], + ) + await asyncio.wait_for(warmed, timeout=2.0) + call_sizes.clear() + + first = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="first")], + ["dense"], + ) + # The process loop is already runnable ahead of this submit. Both + # are ready before this task suspends; no timed arrival is required. + second = await asyncio.create_task(submit_ready()) + await asyncio.wait_for(asyncio.gather(first, second), timeout=2.0) + assert call_sizes == [2] + finally: + await worker.stop() + + @pytest.mark.parametrize("warmup", [False, True]) + def test_ready_producer_callbacks_join_the_idle_head( + self, mock_adapter: MagicMock, warmup: bool, ready_loop_factory + ) -> None: + async def check() -> None: + """Ready transport-like callbacks create request tasks behind the worker. + + A bare sleep(0) resumes the worker before those tasks. The checkpoint + lets their synchronous preparation and submit run without a timed wait. + """ + call_sizes: list[int] = [] + producer_callbacks: list[int] = [] + producers_seen_at_dispatch: list[int] = [] + + def counting_encode(items, output_types, **kwargs): + call_sizes.append(len(items)) + producers_seen_at_dispatch.append(len(producer_callbacks)) + return EncodeOutput(dense=np.array([[0.1, 0.2, 0.3]] * len(items)), batch_size=len(items)) + + mock_adapter.encode.side_effect = counting_encode + worker = ModelWorker( + mock_adapter, + WorkerConfig(max_batch_tokens=1000, max_batch_requests=64, idle_coalesce_ms=3), + ) + await worker.start() + tasks: list[asyncio.Task[None]] = [] + submitted: list[asyncio.Future[WorkerResult]] = [] + all_submitted = asyncio.Event() + + async def submit_tail(index: int) -> None: + submitted.append( + await worker.submit([make_text_item([1, 2], 0)], [Item(text=f"tail {index}")], ["dense"]) + ) + if len(submitted) == 7: + all_submitted.set() + + def produce(index: int) -> None: + producer_callbacks.append(index) + tasks.append(asyncio.create_task(submit_tail(index))) + + try: + if warmup: + warmed = await worker.submit([make_text_item([1, 2], 0)], [Item(text="warmup")], ["dense"]) + await asyncio.wait_for(warmed, timeout=2) + call_sizes.clear() + producers_seen_at_dispatch.clear() + first = await worker.submit([make_text_item([1, 2], 0)], [Item(text="head")], ["dense"]) + loop = asyncio.get_running_loop() + for index in range(7): + loop.call_soon(produce, index) + await asyncio.wait_for(all_submitted.wait(), timeout=2) + await asyncio.wait_for(asyncio.gather(first, *submitted), timeout=2) + assert producer_callbacks == list(range(7)) + assert producers_seen_at_dispatch == [7] + assert call_sizes == [8] + finally: + await asyncio.gather(*tasks, return_exceptions=True) + await worker.stop() + + with asyncio.Runner(loop_factory=ready_loop_factory) as runner: + runner.run(check()) + + @pytest.mark.asyncio + async def test_already_queued_burst_fuses(self, mock_adapter: MagicMock) -> None: + """Several items already pending when an idle worker selects still fuse. + + Queued before the process loop runs, so selection sees the whole + burst rather than a request that arrived alone (#373). + """ + call_sizes: list[int] = [] + + def counting_encode(items, output_types, **kwargs): + call_sizes.append(len(items)) + return EncodeOutput( + dense=np.array([[0.1, 0.2, 0.3]] * len(items)), + batch_size=len(items), + ) + + mock_adapter.encode.side_effect = counting_encode + + config = WorkerConfig( + max_batch_tokens=1000, + max_batch_requests=64, + max_batch_wait_ms=500, + coalesce_ms=400, + coalesce_ratio=1.0, + idle_coalesce_ms=200, + ) + worker = ModelWorker(mock_adapter, config) + await worker.start() + start_loop = worker._ensure_process_loop_started + worker._ensure_process_loop_started = lambda: None # type: ignore[method-assign] + try: futures = [] for i in range(4): @@ -623,12 +864,12 @@ def counting_encode(items, output_types, **kwargs): ["dense"], ) futures.append(future) - # Stagger arrivals well inside the 60ms idle window. - await asyncio.sleep(0.005) + assert worker.pending_count == 4 + del worker._ensure_process_loop_started + start_loop() await asyncio.gather(*futures) - # All four staggered requests fused into a single forward. assert sum(call_sizes) == 4 assert call_sizes == [4], f"expected one fused batch, got: {call_sizes}" @@ -637,11 +878,14 @@ def counting_encode(items, output_types, **kwargs): @pytest.mark.asyncio async def test_staggered_burst_shreds_without_idle_window(self, mock_adapter: MagicMock) -> None: - """Contrast for the test above: ``idle_coalesce_ms=0`` restores the - legacy immediate idle dispatch, and the same staggered burst shreds - into multiple small forwards. This is the mutation the idle window - exists to kill — if the window stops being applied, the fusing test - above degrades into THIS shape and fails. + """``idle_coalesce_ms=0`` does not hold the queued tail the window fuses. + + Same shape as ``test_staggered_burst_fuses_with_idle_window``: the + head runs alone, two requests are already queued before the loop + selects, and one more arrives 10ms later. With the window off that + pair is dispatched immediately, so the later arrival does not join + it (``[1, 2, 1]``). With the window on the same tail is held and + fuses (``[1, 3]``). This is the escape hatch, not the default. """ call_sizes: list[int] = [] @@ -658,38 +902,61 @@ def counting_encode(items, output_types, **kwargs): max_batch_tokens=1000, max_batch_requests=64, max_batch_wait_ms=500, - coalesce_ms=200, - coalesce_ratio=0.5, + coalesce_ms=400, + coalesce_ratio=1.0, idle_coalesce_ms=0, ) worker = ModelWorker(mock_adapter, config) await worker.start() try: - futures = [] - for i in range(4): - future = await worker.submit( - [make_text_item([1, 2], 0)], - [Item(text=f"text {i}")], - ["dense"], - ) - futures.append(future) - await asyncio.sleep(0.005) + started = time.monotonic() + first = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="text 0")], + ["dense"], + ) + await asyncio.wait_for(first, timeout=2.0) + head_ms = (time.monotonic() - started) * 1000 + assert head_ms < 50, f"lone head waited {head_ms:.1f}ms" + assert call_sizes == [1], f"expected the head to run alone, got: {call_sizes}" + + # Same already-queued tail as the fusing test. The window is 0, + # so the pair must leave before the later arrival instead of being held. + second = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="text 1")], + ["dense"], + ) + third = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="text 2")], + ["dense"], + ) + assert worker.pending_count == 2 + await asyncio.sleep(0.01) + assert worker.pending_count == 0, "tail was held instead of dispatched" + fourth = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="text 3")], + ["dense"], + ) - await asyncio.gather(*futures) + await asyncio.gather(second, third, fourth) assert sum(call_sizes) == 4 - # The first arrival dispatches alone (immediate idle dispatch), so - # the burst cannot land in a single forward. - assert len(call_sizes) >= 2, f"expected shredded batches without the window, got: {call_sizes}" + assert call_sizes == [1, 2, 1], f"expected the queued tail to shred, got: {call_sizes}" finally: await worker.stop() @pytest.mark.asyncio - async def test_lone_request_waits_only_the_idle_window(self, mock_adapter: MagicMock) -> None: - """A lone request must not wait the busy-path coalesce/batch windows — - only the small idle accumulation window bounds its dispatch latency. + async def test_lone_request_does_not_wait_the_idle_window(self, mock_adapter: MagicMock) -> None: + """A request that stays alone is not held for ``idle_coalesce_ms`` (#373). + + The busy-path coalesce and batch windows are also not charged. The + bound is loose on purpose: a hard 1ms assert flakes under CI noise, + but it must finish far below the 200ms idle window. """ mock_adapter.encode.return_value = EncodeOutput( dense=np.array([[0.1, 0.2, 0.3]]), @@ -702,7 +969,7 @@ async def test_lone_request_waits_only_the_idle_window(self, mock_adapter: Magic max_batch_wait_ms=500, coalesce_ms=400, coalesce_ratio=1.0, - idle_coalesce_ms=5, + idle_coalesce_ms=200, ) worker = ModelWorker(mock_adapter, config) await worker.start() @@ -717,9 +984,136 @@ async def test_lone_request_waits_only_the_idle_window(self, mock_adapter: Magic await asyncio.wait_for(future, timeout=2.0) elapsed_ms = (time.monotonic() - start) * 1000 - # ~5ms idle window + mock inference; far below the 400ms coalesce - # and 500ms batch windows. - assert elapsed_ms < 200 + assert elapsed_ms < 50, f"lone request waited {elapsed_ms:.1f}ms" + + finally: + await worker.stop() + + @pytest.mark.asyncio + async def test_multi_item_request_does_not_wait_the_idle_window(self, mock_adapter: MagicMock) -> None: + """One submit of two items is one request and skips ``idle_coalesce_ms`` (#373).""" + call_sizes: list[int] = [] + + def counting_encode(items, output_types, **kwargs): + call_sizes.append(len(items)) + return EncodeOutput( + dense=np.array([[0.1, 0.2, 0.3]] * len(items)), + batch_size=len(items), + ) + + mock_adapter.encode.side_effect = counting_encode + + config = WorkerConfig( + max_batch_tokens=1000, + max_batch_requests=64, + max_batch_wait_ms=500, + coalesce_ms=400, + coalesce_ratio=1.0, + idle_coalesce_ms=200, + ) + worker = ModelWorker(mock_adapter, config) + await worker.start() + + try: + start = time.monotonic() + future = await worker.submit( + [make_text_item([1, 2], 0), make_text_item([3, 4], 1)], + [Item(text="hello-0"), Item(text="hello-1")], + ["dense"], + ) + await asyncio.wait_for(future, timeout=2.0) + elapsed_ms = (time.monotonic() - start) * 1000 + + assert elapsed_ms < 50, f"multi-item request waited {elapsed_ms:.1f}ms" + assert call_sizes == [2] + + finally: + await worker.stop() + + @pytest.mark.asyncio + async def test_sequential_request_does_not_wait_the_idle_window(self, mock_adapter: MagicMock) -> None: + """A finished request does not make the next submit wait (#373). + + The second request is sent as soon as the first has been awaited, well + inside ``idle_coalesce_ms``. It is not a tail: nothing is still pending + on this batcher. + """ + mock_adapter.encode.return_value = EncodeOutput( + dense=np.array([[0.1, 0.2, 0.3]]), + batch_size=1, + ) + + config = WorkerConfig( + max_batch_tokens=1000, + max_batch_requests=64, + max_batch_wait_ms=500, + coalesce_ms=400, + coalesce_ratio=1.0, + idle_coalesce_ms=200, + ) + worker = ModelWorker(mock_adapter, config) + await worker.start() + + try: + first = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="first")], + ["dense"], + ) + await asyncio.wait_for(first, timeout=2.0) + + start = time.monotonic() + second = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="second")], + ["dense"], + ) + await asyncio.wait_for(second, timeout=2.0) + elapsed_ms = (time.monotonic() - start) * 1000 + + assert elapsed_ms < 50, f"sequential request waited {elapsed_ms:.1f}ms" + + finally: + await worker.stop() + + @pytest.mark.asyncio + async def test_other_lora_arrival_does_not_hold_this_batcher(self, mock_adapter: MagicMock) -> None: + """A different LoRA's arrival does not make this batcher wait (#373).""" + mock_adapter.encode.return_value = EncodeOutput( + dense=np.array([[0.1, 0.2, 0.3]]), + batch_size=1, + ) + + config = WorkerConfig( + max_batch_tokens=1000, + max_batch_requests=64, + max_batch_wait_ms=500, + coalesce_ms=400, + coalesce_ratio=1.0, + idle_coalesce_ms=200, + ) + worker = ModelWorker(mock_adapter, config) + await worker.start() + + try: + other = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="lora")], + ["dense"], + options={"lora": "adapter-a"}, + ) + await asyncio.wait_for(other, timeout=2.0) + + start = time.monotonic() + base = await worker.submit( + [make_text_item([1, 2], 0)], + [Item(text="base")], + ["dense"], + ) + await asyncio.wait_for(base, timeout=2.0) + elapsed_ms = (time.monotonic() - start) * 1000 + + assert elapsed_ms < 50, f"other LoRA held this batcher for {elapsed_ms:.1f}ms" finally: await worker.stop() diff --git a/packages/sie_server/tests/core/test_worker_drain.py b/packages/sie_server/tests/core/test_worker_drain.py index 320b0aca1..c9c109ba3 100644 --- a/packages/sie_server/tests/core/test_worker_drain.py +++ b/packages/sie_server/tests/core/test_worker_drain.py @@ -46,9 +46,10 @@ def _parked_config() -> WorkerConfig: """Config whose batcher never yields a batch during a test. Every dispatch trigger (cost limit, request limit, first-request timeout, - coalesce window, idle accumulation window) is pushed far out of reach, so - submitted items stay queued in the ``BatchFormer`` for the whole test and - the drain is the only thing that can complete them. This is the state a + coalesce window, idle tail window) is pushed far out of reach. A lone + idle request is still dispatched immediately (#373), so ``_submit`` queues + a second request on the same batcher and the tail window holds the burst. + The drain is the only thing that can complete it. This is the state a real worker is in whenever it is evicted with a backlog. """ return WorkerConfig( @@ -122,19 +123,38 @@ def _item(index: int = 0) -> TextPreparedItem: return make_text_item([1, 2, 3, 4, 5], index) +def _retrieve_exception(future: asyncio.Future[Any]) -> None: + """Retrieve a parked partner's error so a drained future does not warn.""" + if not future.cancelled(): + future.exception() + + async def _submit( worker: ModelWorker, *, count: int = 1, options: dict[str, Any] | None = None, ) -> asyncio.Future[WorkerResult]: - """Queue one request of ``count`` items and let the process loop park.""" + """Queue one request of ``count`` items and let the process loop park. + + A lone idle request, including one submit of several items, is dispatched + without waiting ``idle_coalesce_ms`` (#373). A second request on the same + batcher is submitted through the public API before the loop selects, so + the burst is a tail and the parked window holds both inside ``get_batch``. + """ future = await worker.submit( [_item(i) for i in range(count)], [Item(text=f"hello-{i}") for i in range(count)], ["dense"], options=options, ) + partner = await worker.submit( + [_item(count)], + [Item(text="park")], + ["dense"], + options=options, + ) + partner.add_done_callback(_retrieve_exception) # Give the process loop a chance to start and block on the batcher, so # the drain is exercised against a live (not merely un-started) loop. await asyncio.sleep(0) @@ -149,7 +169,8 @@ async def test_queued_request_fails_with_retryable_error(self, mock_adapter: Mag worker = ModelWorker(mock_adapter, _parked_config(), model_name="test-model") await worker.start() future = await _submit(worker) - assert worker.pending_count == 1 + # The request under test plus the one-item request that holds the window. + assert worker.pending_count == 2 assert not future.done() await worker.stop() @@ -170,7 +191,8 @@ async def test_no_item_is_left_pending_in_any_batcher(self, mock_adapter: MagicM lora_a = await _submit(worker, options={"lora": "adapter-a"}) lora_b = await _submit(worker, options={"lora": "adapter-b"}) assert set(worker._batchers) == {None, "adapter-a", "adapter-b"} - assert worker.pending_count == 3 + # Each batcher has the request under test and the request that holds it. + assert worker.pending_count == 6 await worker.stop() @@ -191,7 +213,8 @@ async def test_multi_item_request_fails_its_single_future_once(self, mock_adapte worker = ModelWorker(mock_adapter, _parked_config(), model_name="test-model") await worker.start() future = await _submit(worker, count=4) - assert worker.pending_count == 4 + # Four items share the future under test; one more item holds the window. + assert worker.pending_count == 5 await worker.stop() diff --git a/packages/sie_server/tests/core/test_worker_idle_wakeup.py b/packages/sie_server/tests/core/test_worker_idle_wakeup.py index 1b0d2d42c..1d0960060 100644 --- a/packages/sie_server/tests/core/test_worker_idle_wakeup.py +++ b/packages/sie_server/tests/core/test_worker_idle_wakeup.py @@ -158,43 +158,49 @@ async def observed_wait(): return await original_wait() monkeypatch.setattr(batcher._batch_ready, "wait", observed_wait) - future = await _submit(worker, "queued") + # Multiple pending requests use the idle coalescing window; a lone + # request deliberately bypasses that window after the ready checkpoint. + futures = [await _submit(worker, "queued-head"), await _submit(worker, "queued-tail")] await asyncio.wait_for(waiting.wait(), _TIMEOUT) await asyncio.wait_for(worker.stop(), _TIMEOUT) - with pytest.raises(WorkerDrainedError): - await future + for future in futures: + with pytest.raises(WorkerDrainedError): + await future assert worker.pending_count == 0 assert worker._process_task is None worker.adapter.encode.assert_not_called() @pytest.mark.parametrize( - ("case", "steps", "expected_time", "expected_count"), + ("case", "initial_count", "steps", "expected_time", "expected_count"), [ - ("lone", [(2.0, False), (3.1, False)], 3.1, 1), - ("burst", [(2.0, True), (4.0, False), (5.1, False)], 5.1, 2), - ("busy", [(3.1, False), (7.6, False)], 7.6, 1), - ("full", [(2.0, True)], 2.0, 2), - ("timeout", [(2.0 * i, True) for i in range(1, 8)] + [(15.1, False)], 15.1, 8), - ("zero", [], 0.0, 1), + ("lone", 1, [], 0.0, 1), + ("burst", 2, [(2.0, True), (4.0, False), (5.1, False)], 5.1, 3), + ("busy", 1, [(3.1, False), (7.6, False)], 7.6, 1), + ("full", 2, [(2.0, True)], 2.0, 3), + ("timeout", 2, [(2.0 * i, True) for i in range(1, 8)] + [(15.1, False)], 15.1, 9), + ("zero", 1, [], 0.0, 1), ], ) -async def test_original_coalescing_deadlines_and_membership( - worker, monkeypatch, case, steps, expected_time, expected_count +async def test_coalescing_deadlines_and_lone_dispatch_membership( + worker, monkeypatch, case, initial_count, steps, expected_time, expected_count ): """Use a controlled batcher clock, not fragile wall-time upper bounds. - The real selector and BatchFormer must retain the default 3ms idle quiet - period, its reset on arrival, the 15ms first-item deadline, fullness and - the uncapped 7.5ms busy quiet period. No model or thread timing is measured. + The real selector bypasses the timer for a lone request. Multiple pending + requests retain the 3ms idle quiet period and its reset on arrival, the + 15ms first-item deadline, fullness and the uncapped 7.5ms busy quiet period. + No model or thread timing is measured. """ worker.config.idle_coalesce_ms = 0.0 if case == "zero" else 3.0 if case == "full": - worker._batch_config.max_batch_requests = 2 + worker._batch_config.max_batch_requests = 3 now = 0.0 monkeypatch.setattr(batcher_module, "time", SimpleNamespace(monotonic=lambda: now / 1000)) monkeypatch.setattr(worker, "_ensure_process_loop_started", lambda: None) futures = [await _submit(worker, "first")] + for index in range(1, initial_count): + futures.append(await _submit(worker, f"item-{index}")) remaining = iter(steps) observed = [] diff --git a/tools/benchmarks/idle_worker/.gitignore b/tools/benchmarks/idle_worker/.gitignore new file mode 100644 index 000000000..c36187138 --- /dev/null +++ b/tools/benchmarks/idle_worker/.gitignore @@ -0,0 +1,2 @@ +_inputs/ +results/ diff --git a/tools/benchmarks/idle_worker/README.md b/tools/benchmarks/idle_worker/README.md new file mode 100644 index 000000000..cc9cbce2f --- /dev/null +++ b/tools/benchmarks/idle_worker/README.md @@ -0,0 +1,62 @@ +# Idle-worker scheduling comparison + +This benchmark compares two exact SIE source revisions on one NVIDIA GPU. It +uses `sie_sdk.SIEClient` against the local Python worker HTTP endpoint, records +the actual `X-Queue-Time` headers, and observes worker batch sizes. It targets +the Python worker scheduler changed in [PR #658](https://github.com/superlinked/sie/pull/658). +It does not measure the queue-only gateway/sidecar path, whose preformed batches +have a different scheduling owner. + +The default workload uses the catalog-pinned `knowledgator/gliclass-large-v1.0` +checkpoint, four synthetic support requests, and five fixed classification +labels. Both arms use their checked-in default model profile and scheduling +configuration. The archive includes the server source, model catalog and bundle +definitions. Preparation refuses different dependency locks. + +From an initialized checkout: + +```sh +mise exec -- python tools/benchmarks/idle_worker/prepare.py \ + --baseline b0c26b02d21dbd146646c41b2b4f42627a0c9ebd \ + --candidate a7c024f6f53cce0f8e43825845c5ad6965c9c792 +``` + +On the GPU machine, with the same locked workspace and prepared `_inputs/`: + +```sh +mise exec -- python tools/benchmarks/idle_worker/run.py \ + --output /tmp/sie-idle-worker-results +``` + +The default order is baseline/candidate/candidate/baseline/baseline/candidate. +Each fresh server warms request concurrency 1, 8 and 32, then measures: + +- 100 sequential requests, including every queue-time sample; +- four simultaneous HTTP bursts each at concurrency 8 and 32; +- 12 seconds of closed-loop traffic each at concurrency 8 and 32, including + the final outstanding requests in elapsed time. + +HTTP bursts can reach the worker at different times: distinguish an already +queued burst from the PR's intended first-request/tail split. The batch trace +records item IDs so the observed grouping is inspectable. A wrapper records a +batch only after normal dispatch completes; it does not replace scheduling or +inference. The existing worker unit tests cover deterministic queued-burst +formation. + +Each completed phase is written before the next starts. If a later phase fails, +the trial is marked `failed` and its completed measurements remain available. +Trial JSON is replaced atomically so an incomplete write preserves the last +completed phase. Batch records are flushed as JSONL after each dispatch, so +SIGTERM during server shutdown cannot +discard them. If a trace is missing, request timings remain available and the +trial is marked `timing_complete_trace_missing`. +Reported p95 values use nearest rank: sort the samples and select rank +`ceil(0.95 * n)`, counting from one. + +Use an otherwise idle GPU and reserve enough CPU for the client and server. +The server uses four Torch CPU threads. Keep the GPU allocation, dependencies, +input texts and warmup identical for both arms; compare all repeated trials +before interpreting a small throughput change. The output includes source and +archive hashes, runtime versions, GPU identity, raw sequential headers, HTTP +latencies, batch traces, and each server's log. Generated snapshots and results +are ignored by Git. diff --git a/tools/benchmarks/idle_worker/__init__.py b/tools/benchmarks/idle_worker/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tools/benchmarks/idle_worker/prepare.py b/tools/benchmarks/idle_worker/prepare.py new file mode 100644 index 000000000..af2b71970 --- /dev/null +++ b/tools/benchmarks/idle_worker/prepare.py @@ -0,0 +1,51 @@ +"""Snapshot two Git revisions for a same-machine idle-worker comparison.""" + +import argparse +import hashlib +import json +import shutil +import subprocess +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[3] +INPUTS = Path(__file__).parent / "_inputs" +PATHS = ("packages/sie_server/src", "packages/sie_server/models", "packages/sie_server/bundles") +GIT = shutil.which("git") + + +def git(*args: str) -> bytes: + if GIT is None: + raise RuntimeError("git is required") + return subprocess.check_output([GIT, *args], cwd=ROOT) # noqa: S603 — local Git revisions, no shell. + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline", required=True) + parser.add_argument("--candidate", required=True) + args = parser.parse_args() + INPUTS.mkdir(exist_ok=True) + manifest = {} + locks = [] + for arm, revision in (("baseline", args.baseline), ("candidate", args.candidate)): + commit = git("rev-parse", f"{revision}^{{commit}}").decode().strip() + archive = git("archive", "--format=tar", commit, *PATHS) + (INPUTS / f"{arm}.tar").write_bytes(archive) + lock = git("show", f"{commit}:uv.lock") + locks.append(lock) + manifest[arm] = { + "commit": commit, + "archive_sha256": hashlib.sha256(archive).hexdigest(), + "uv_lock_sha256": hashlib.sha256(lock).hexdigest(), + } + if locks[0] != locks[1] or locks[1] != (ROOT / "uv.lock").read_bytes(): + raise SystemExit("Both revisions and the execution workspace must use the same uv.lock") + manifest["source_diff"] = git( + "diff", "--stat", manifest["baseline"]["commit"], manifest["candidate"]["commit"], "--", *PATHS + ).decode() + (INPUTS / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/benchmarks/idle_worker/run.py b/tools/benchmarks/idle_worker/run.py new file mode 100644 index 000000000..ebd3bf9ab --- /dev/null +++ b/tools/benchmarks/idle_worker/run.py @@ -0,0 +1,302 @@ +"""Compare idle-worker HTTP latency and throughput on one GPU via SIEClient.""" + +import argparse +import hashlib +import importlib.metadata +import itertools +import json +import math +import os +import platform +import shutil +import socket +import statistics +import subprocess +import sys +import tarfile +import tempfile +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from http import HTTPStatus +from pathlib import Path + +import httpx +from sie_sdk import SIEClient +from sie_sdk.types import Item + +HERE = Path(__file__).resolve().parent +ROOT = HERE.parents[2] +LABELS = ["billing", "technical support", "account access", "sales", "other"] +TEXTS = [ + "The payment for my annual subscription was charged twice yesterday. Please check the invoice and refund the duplicate charge. I need the account to stay active while your team reviews the transaction.", + "Our data export has stopped working after the latest update. The dashboard shows the job as complete, but the downloaded file contains only the column headings. We tried again from another browser with the same result. Please help us restore the missing records before the report is due tomorrow.", + "I cannot log in to my account because the verification message never arrives. I checked the spam folder and confirmed that the email address is correct. My colleagues can still access their accounts. Please let me know how to recover access without creating a second account.", + "We are evaluating the service for a team of twenty analysts and would like information about the available subscription plans. We expect to process several thousand documents each week. Could you explain how usage is measured and whether the service supports a shared workspace for our team? We would also like to schedule a short demonstration next week.", +] + + +def distribution(values: list[float]) -> dict: + """Summarize samples using the nearest-rank definition of p95.""" + ordered = sorted(values) + return { + "n": len(values), + "median": statistics.median(values), + "p95": ordered[math.ceil(len(values) * 0.95) - 1], + "max": max(values), + } + + +def timing_sample(response: httpx.Response) -> dict[str, float] | None: + if not response.request.url.path.startswith("/v1/extract/") or response.status_code != HTTPStatus.OK: + return None + return { + key: float(response.headers[header]) + for key, header in ( + ("queue_ms", "X-Queue-Time"), + ("inference_ms", "X-Inference-Time"), + ("total_ms", "X-Total-Time"), + ) + } + + +def measure( + client: SIEClient, model: str, phase: str, concurrency: int, *, count: int = 0, duration: float = 0 +) -> dict: + barrier = threading.Barrier(concurrency) + counter = itertools.count() + start = 0.0 + latencies: list[float] = [] + + def lane() -> None: + barrier.wait() + while True: + index = next(counter) + if (count and index >= count) or (duration and time.perf_counter() - start >= duration): + return + before = time.perf_counter() + result = client.extract( + model, + Item(text=TEXTS[index % len(TEXTS)], id=f"{phase}:{index}"), + labels=LABELS, + wait_for_capacity=False, + max_oom_retries=0, + ) + if result.get("error"): + raise RuntimeError(result["error"]) + if not result.get("classifications"): + raise RuntimeError("No classifications returned") + latencies.append((time.perf_counter() - before) * 1000) + + with ThreadPoolExecutor(max_workers=concurrency) as executor: + start = time.perf_counter() + futures = [executor.submit(lane) for _ in range(concurrency)] + for future in futures: + future.result() + elapsed = time.perf_counter() - start + return { + "phase": phase, + "concurrency": concurrency, + "elapsed_s": elapsed, + "requests": len(latencies), + "requests_per_s": len(latencies) / elapsed, + "latency_ms": distribution(latencies), + } + + +def wait_for_server(process: subprocess.Popen, port: int) -> None: + deadline = time.monotonic() + 300 + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError(f"Server exited with {process.returncode}") + try: + with socket.create_connection(("127.0.0.1", port), timeout=1): + return + except OSError: + time.sleep(0.25) + raise TimeoutError("Server did not bind within 300 seconds") + + +def trial(args: argparse.Namespace, arm: str, trial_index: int, source: Path, output: Path) -> dict: + name = f"{trial_index:02d}-{arm}" + samples: list[dict] = [] + + def capture(response: httpx.Response) -> None: + if (sample := timing_sample(response)) is not None: + samples.append(sample) + + environment = os.environ.copy() + environment.update( + { + "PYTHONPATH": str(source / "packages/sie_server/src") + os.pathsep + environment.get("PYTHONPATH", ""), + "SIE_BUNDLE": "default", + "SIE_TELEMETRY_DISABLED": "true", + "SIE_TORCH_NUM_THREADS": "4", + "TOKENIZERS_PARALLELISM": "false", + "PYTHONUNBUFFERED": "1", + "SIE_LOG_LEVEL": "warning", + } + ) + trace = output / f"{name}-batches.jsonl" + argv = [ + sys.executable, + str(HERE / "server.py"), + "--model", + args.model, + "--models-dir", + str(source / "packages/sie_server/models"), + "--trace", + str(trace), + "--port", + str(args.port), + "--device", + args.device, + ] + phases = [] + result = {"arm": arm, "trial": trial_index, "phases": phases, "status": "running"} + result_path = output / f"{name}.json" + + def persist_result() -> None: + pending = result_path.with_suffix(".json.tmp") + pending.write_text(json.dumps(result, indent=2) + "\n") + pending.replace(result_path) + + def record_phase(phase: dict) -> None: + phases.append(phase) + persist_result() + + with (output / f"{name}-server.log").open("w") as log: + process = subprocess.Popen(argv, env=environment, stdout=log, stderr=subprocess.STDOUT) # noqa: S603 — fixed local server, no shell. + try: + wait_for_server(process, args.port) + base_url = f"http://127.0.0.1:{args.port}" + transport = httpx.Client( + base_url=base_url, + timeout=300, + event_hooks={"response": [capture]}, + limits=httpx.Limits(max_connections=64, max_keepalive_connections=64), + ) + with SIEClient(base_url, timeout_s=300, http_client=transport) as client: + for concurrency in (1, 8, 32): + measure(client, args.model, f"warmup{concurrency}", concurrency, count=64) + samples.clear() + sequential = measure(client, args.model, "sequential", 1, count=args.sequential) + if len(samples) != args.sequential: + raise RuntimeError(f"Expected {args.sequential} successful timing samples, got {len(samples)}") + sequential["queue_ms"] = distribution([sample["queue_ms"] for sample in samples]) + sequential["inference_ms"] = distribution([sample["inference_ms"] for sample in samples]) + sequential["queue_under_1ms"] = sum(sample["queue_ms"] < 1 for sample in samples) + sequential["samples"] = samples.copy() + record_phase(sequential) + for concurrency in (8, 32): + for burst in range(4): + time.sleep(0.05) + record_phase( + measure(client, args.model, f"burst{concurrency}-{burst}", concurrency, count=concurrency) + ) + measure(client, args.model, f"warmthroughput{concurrency}", concurrency, count=128) + record_phase( + measure(client, args.model, f"throughput{concurrency}", concurrency, duration=args.duration) + ) + except BaseException as error: + result["status"] = "failed" + try: + persist_result() + except OSError as persistence_error: + error.add_note(f"Could not persist failed trial: {persistence_error}") + raise + finally: + process.terminate() + try: + process.wait(timeout=40) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + result.update(status="complete", server_returncode=process.returncode) + if trace.exists(): + batches = [json.loads(line) for line in trace.read_text().splitlines()] + for phase in phases: + prefix = phase["phase"] + ":" + phase["worker_batch_sizes"] = [ + record["size"] for record in batches if any(i.startswith(prefix) for i in record["request_ids"]) + ] + else: + result["status"] = "timing_complete_trace_missing" + persist_result() + return result + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default="knowledgator/gliclass-large-v1.0") + parser.add_argument("--output", type=Path, default=HERE / "results") + parser.add_argument("--duration", type=float, default=12) + parser.add_argument("--sequential", type=int, default=100) + parser.add_argument("--order", default="baseline,candidate,candidate,baseline,baseline,candidate") + parser.add_argument("--port", type=int, default=8373) + parser.add_argument("--device", default="cuda") + args = parser.parse_args() + if args.duration <= 0 or args.sequential < 1: + parser.error("duration and sequential must be positive") + order = args.order.split(",") + if not order or any(arm not in {"baseline", "candidate"} for arm in order): + parser.error("order must contain baseline/candidate arms") + args.output.mkdir(parents=True, exist_ok=True) + manifest = json.loads((HERE / "_inputs/manifest.json").read_text()) + nvidia_smi = shutil.which("nvidia-smi") + if nvidia_smi is None: + raise RuntimeError("nvidia-smi is required") + environment = { + "python": platform.python_version(), + "packages": { + name: importlib.metadata.version(name) + for name in ("torch", "transformers", "gliclass", "gliformer", "httpx", "uvicorn") + }, + "nvidia_smi": subprocess.check_output( # noqa: S603 — fixed GPU metadata query, no shell. + [nvidia_smi, "--query-gpu=name,uuid,driver_version,memory.total", "--format=csv,noheader"], text=True + ).strip(), + "cpu_count": os.cpu_count(), + "torch_threads": 4, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES", ""), + } + result = { + "manifest": manifest, + "environment": environment, + "model": args.model, + "texts": TEXTS, + "labels": LABELS, + "trials": [], + } + with tempfile.TemporaryDirectory(prefix="sie-idle-worker-") as temporary: + roots = {} + for arm in set(order): + archive = HERE / f"_inputs/{arm}.tar" + if hashlib.sha256(archive.read_bytes()).hexdigest() != manifest[arm]["archive_sha256"]: + raise RuntimeError(f"{arm} archive hash mismatch") + roots[arm] = Path(temporary) / arm + roots[arm].mkdir() + with tarfile.open(archive) as snapshot: + snapshot.extractall(roots[arm], filter="data") + for index, arm in enumerate(order): + print(f"IDLE_WORKER_START trial={index} arm={arm}", flush=True) + current = trial(args, arm, index, roots[arm], args.output) + result["trials"].append(current) + (args.output / "results.json").write_text(json.dumps(result, indent=2) + "\n") + print( + json.dumps( + { + "arm": arm, + "trial": index, + "phases": [ + {k: v for k, v in p.items() if k not in {"samples", "worker_batch_sizes"}} + for p in current["phases"] + ], + } + ), + flush=True, + ) + print("IDLE_WORKER_COMPLETE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/benchmarks/idle_worker/server.py b/tools/benchmarks/idle_worker/server.py new file mode 100644 index 000000000..68f3cfaf7 --- /dev/null +++ b/tools/benchmarks/idle_worker/server.py @@ -0,0 +1,50 @@ +"""Local worker HTTP server with observational batch tracing for the benchmark.""" + +import argparse +import hashlib +import json +from pathlib import Path + +import uvicorn +from sie_server.app.app_factory import AppFactory +from sie_server.app.app_state_config import AppStateConfig +from sie_server.core.batcher import FormattedBatch, HasCost +from sie_server.core.worker import ModelWorker, model_worker +from sie_server.core.worker.types import RequestMetadata + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("--model", required=True) +parser.add_argument("--models-dir", type=Path, required=True) +parser.add_argument("--trace", type=Path, required=True) +parser.add_argument("--port", type=int, default=8373) +parser.add_argument("--device", default="cuda") +args = parser.parse_args() +trace_stream = args.trace.open("w") +original_process_batch = ModelWorker._process_batch + + +async def traced_process_batch( + self: ModelWorker, batch: FormattedBatch[HasCost, RequestMetadata], *, engine_queue_owned: bool = False +) -> None: + await original_process_batch(self, batch, engine_queue_owned=engine_queue_owned) + ids = sorted({str(item.id) for metadata in batch.metadata for item in metadata.items}) + trace_stream.write(json.dumps({"size": batch.size, "request_ids": ids}) + "\n") + trace_stream.flush() + + +ModelWorker._process_batch = traced_process_batch # ty: ignore[invalid-assignment] +source = Path(model_worker.__file__) +print( + json.dumps({"worker_path": str(source), "worker_sha256": hashlib.sha256(source.read_bytes()).hexdigest()}), + flush=True, +) +app = AppFactory.create_app( + AppStateConfig( + models_dir=args.models_dir, + device=args.device, + model_filter=[args.model], + preload_models=[args.model], + remote_serving=False, + ) +) +uvicorn.run(app, host="127.0.0.1", port=args.port, log_level="warning", access_log=False, loop="asyncio") diff --git a/tools/benchmarks/idle_worker/tests/test_prepare.py b/tools/benchmarks/idle_worker/tests/test_prepare.py new file mode 100644 index 000000000..f5a01add2 --- /dev/null +++ b/tools/benchmarks/idle_worker/tests/test_prepare.py @@ -0,0 +1,59 @@ +"""Snapshot provenance must stay stable even if an input branch moves.""" + +import json +import sys +from pathlib import Path + +import pytest + +from tools.benchmarks.idle_worker import prepare + + +def test_diff_describes_snapshots_when_input_branch_moves(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(prepare, "ROOT", tmp_path) + monkeypatch.setattr(prepare, "INPUTS", tmp_path / "_inputs") + git = prepare.git + git("init", "--quiet") + (tmp_path / "uv.lock").write_text("same lock\n") + for directory in prepare.PATHS: + target = tmp_path / directory + target.mkdir(parents=True) + (target / "probe.txt").write_text("baseline\n") + + def commit() -> str: + git("add", ".") + git( + "-c", + "user.name=Benchmark Test", + "-c", + "user.email=benchmark@example.test", + "commit", + "--quiet", + "-m", + "snapshot", + ) + return git("rev-parse", "HEAD").decode().strip() + + baseline = commit() + (tmp_path / prepare.PATHS[0] / "probe.txt").write_text("candidate\n") + candidate = commit() + git("branch", "baseline", baseline) + git("branch", "candidate", candidate) + expected_diff = git("diff", "--stat", baseline, candidate, "--", *prepare.PATHS).decode() + (tmp_path / prepare.PATHS[0] / "later.txt").write_text("not in the archived candidate\n") + moved_candidate = commit() + + def git_with_moving_branch(*args: str) -> bytes: + if args[0] == "diff": + git("update-ref", "refs/heads/candidate", moved_candidate) + return git(*args) + + monkeypatch.setattr(prepare, "git", git_with_moving_branch) + monkeypatch.setattr(sys, "argv", ["prepare.py", "--baseline", "baseline", "--candidate", "candidate"]) + prepare.main() + + manifest = json.loads((prepare.INPUTS / "manifest.json").read_text()) + assert manifest["baseline"]["commit"] == baseline + assert manifest["candidate"]["commit"] == candidate + assert git("rev-parse", "candidate").decode().strip() == moved_candidate + assert manifest["source_diff"] == expected_diff diff --git a/tools/benchmarks/idle_worker/tests/test_run.py b/tools/benchmarks/idle_worker/tests/test_run.py new file mode 100644 index 000000000..332577212 --- /dev/null +++ b/tools/benchmarks/idle_worker/tests/test_run.py @@ -0,0 +1,133 @@ +"""The timing probe follows the SDK's real extraction transport path.""" + +import argparse +import json +from pathlib import Path + +import httpx +import pytest +from sie_sdk import SIEClient +from sie_sdk._msgpack import packb +from sie_sdk.types import Item + +from tools.benchmarks.idle_worker import run +from tools.benchmarks.idle_worker.run import distribution, timing_sample + + +def test_timing_hook_follows_sdk_extract_route() -> None: + samples = [] + + def response_for(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + request=request, + headers={ + "X-Queue-Time": "3.81", + "X-Inference-Time": "20.0", + "X-Total-Time": "24.0", + "Content-Type": "application/msgpack", + }, + content=packb( + { + "model": "knowledgator/gliclass-large-v1.0", + "items": [{"id": "probe", "entities": [], "classifications": [{"label": "billing", "score": 0.9}]}], + } + ), + ) + + transport = httpx.Client( + base_url="http://127.0.0.1:8373", + transport=httpx.MockTransport(response_for), + event_hooks={"response": [lambda response: samples.append(timing_sample(response))]}, + ) + with SIEClient("http://127.0.0.1:8373", http_client=transport) as client: + client.extract("knowledgator/gliclass-large-v1.0", Item(text="Refund please", id="probe"), labels=["billing"]) + assert samples == [{"queue_ms": 3.81, "inference_ms": 20.0, "total_ms": 24.0}] + + +@pytest.mark.parametrize(("count", "p95"), [(1, 1), (20, 19), (100, 95), (101, 96), (300, 285)]) +def test_distribution_uses_nearest_rank_p95(count: int, p95: int) -> None: + assert distribution(list(reversed(range(1, count + 1))))["p95"] == p95 + + +@pytest.mark.parametrize("failure_type", [RuntimeError, KeyboardInterrupt]) +@pytest.mark.parametrize("persistence_fails", [False, True]) +def test_failed_trial_preserves_completed_phases_and_stops_server( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, failure_type: type[BaseException], persistence_fails: bool +) -> None: + stopped = [] + + class Process: + returncode = None + + def terminate(self) -> None: + stopped.append("terminate") + + def wait(self, timeout: int) -> None: + stopped.append(timeout) + + monkeypatch.setattr(run.subprocess, "Popen", lambda *args, **kwargs: Process()) + monkeypatch.setattr(run, "wait_for_server", lambda *args: None) + real_client = httpx.Client + + def response_for(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + request=request, + headers={ + "X-Queue-Time": "0.12", + "X-Inference-Time": "20", + "X-Total-Time": "21", + "Content-Type": "application/msgpack", + }, + content=packb( + { + "model": "probe", + "items": [ + {"id": "sequential:0", "entities": [], "classifications": [{"label": "billing", "score": 0.9}]} + ], + } + ), + ) + + monkeypatch.setattr( + run.httpx, + "Client", + lambda *args, **kwargs: real_client(*args, **kwargs, transport=httpx.MockTransport(response_for)), + ) + failure = failure_type("burst failed after sequential completed") + trace = tmp_path / "00-candidate-batches.jsonl" + trace_record = '{"request_ids": ["sequential:0"], "size": 1}\n' + + def measure(client: SIEClient, model: str, phase: str, concurrency: int, **kwargs) -> dict: + if phase.startswith("warmup"): + return {} + if phase == "sequential": + client.extract(model, Item(text="Refund please", id="sequential:0"), labels=["billing"]) + trace.write_text(trace_record) + return {"phase": phase, "requests": 1} + raise failure + + monkeypatch.setattr(run, "measure", measure) + real_write_text = Path.write_text + + def write_text(path: Path, data: str, *args, **kwargs) -> int: + if persistence_fails and '"status": "failed"' in data: + real_write_text(path, "incomplete") + raise OSError("disk unavailable while recording failure") + return real_write_text(path, data, *args, **kwargs) + + monkeypatch.setattr(Path, "write_text", write_text) + args = argparse.Namespace(model="probe", port=8373, device="cpu", sequential=1, duration=0.01) + with pytest.raises(failure_type) as caught: + run.trial(args, "candidate", 0, tmp_path, tmp_path) + + assert caught.value is failure + assert stopped == ["terminate", 40] + result = json.loads((tmp_path / "00-candidate.json").read_text()) + assert result["status"] == ("running" if persistence_fails else "failed") + assert result["phases"][0]["phase"] == "sequential" + assert result["phases"][0]["samples"] == [{"queue_ms": 0.12, "inference_ms": 20.0, "total_ms": 21.0}] + assert trace.read_text() == trace_record + if persistence_fails: + assert any("disk unavailable" in note for note in caught.value.__notes__)