From 45794f69402d1c334c1238cfc13082cc0596e08f Mon Sep 17 00:00:00 2001 From: svonava Date: Sat, 17 Oct 2026 23:54:25 -0700 Subject: [PATCH] fix(batching): guarantee progress for bypassed requests --- .../sie_server/src/sie_server/core/batcher.py | 63 ++++- .../sie_server/core/worker/model_worker.py | 9 +- .../sie_server/tests/core/test_batcher.py | 227 +++++++++++++++++- 3 files changed, 281 insertions(+), 18 deletions(-) diff --git a/packages/sie_server/src/sie_server/core/batcher.py b/packages/sie_server/src/sie_server/core/batcher.py index c4a5cf4ca..966e57c93 100644 --- a/packages/sie_server/src/sie_server/core/batcher.py +++ b/packages/sie_server/src/sie_server/core/batcher.py @@ -25,8 +25,14 @@ a long item's batch, and waits for at most one flagged item between two batches of its lane. A flagged item with k flagged items ahead of it waits for at most k + 1 batches of batchable items, however many keep arriving. -When nothing pending is flagged, batch formation is exactly the cost-sorted -greedy packing above. +Within the batchable lane, an ordinary cost-sorted batch that bypasses the +oldest item reserves that exact item for the next batchable turn. That turn +packs from the reserved item's cost-sorted suffix, without wrapping to cheaper +items, then ordinary packing resumes. Future cheap arrivals therefore cannot +starve an older item, including an item that exceeds the current cost budget. +An item with k older batchable items waits for at most 2(k + 1) batchable turns; +flagged turns can interleave as above. This bounds dispatch opportunities, not +inference duration. Ordinary batches that include the oldest item are unchanged. Cost semantics vary by modality (modality-native units): - Text: cost = token count @@ -86,6 +92,7 @@ class PendingRequest[I: HasCost, T]: metadata: T arrival_time: float = field(default_factory=time.monotonic) runs_alone: bool = False + _arrival_order: int = field(default=0, init=False, repr=False, compare=False) @dataclass(slots=True) @@ -229,6 +236,8 @@ def __init__(self, config: BatchConfig | None = None) -> None: self._first_request_time: float | None = None self._last_submit_time: float | None = None self._last_batch_alone: bool | None = None + self._next_arrival_order = 0 + self._reserved_request: PendingRequest[I, T] | None = None @property def config(self) -> BatchConfig: @@ -323,6 +332,8 @@ async def submit_many(self, items: list[tuple[I, T]]) -> None: def _append_item(self, item: I, metadata: T) -> None: """Append item to pending list (caller must hold lock).""" request = PendingRequest(item=item, metadata=metadata, runs_alone=getattr(item, "runs_alone", False) is True) + request._arrival_order = self._next_arrival_order + self._next_arrival_order += 1 self._pending.append(request) self._total_cost += item.cost @@ -449,7 +460,10 @@ def _extract_batch( 2. Take requests until max_batch_cost would be exceeded 3. Keep remaining requests for next get_batch() call - This minimizes padding waste by grouping similar-cost sequences. + This minimizes padding waste by grouping similar-cost sequences. If + ordinary packing bypassed the oldest batchable item last time, this + batchable turn starts at that exact reserved item and cannot wrap into + its cheaper prefix. The following batchable turn uses ordinary packing. When a pending request is flagged ``runs_alone``, the order comes from ``_order_runs_alone_lane`` instead, and a flagged request is always taken on its own. @@ -457,9 +471,9 @@ def _extract_batch( Args: max_items: Optional hard cap on the number of items taken, on top of the cost/request limits. Used by ``try_drain`` to bound a - continuous-batch drain to a caller-held snapshot budget so - post-snapshot arrivals are not swept in ahead of the next FCFS - selection. See #1606. + continuous-batch drain to a caller-held snapshot count before + the next FCFS selection. This bounds how many items are taken, + not which arrivals are eligible. See #1606. coalesce_cap_ms: The coalesce cap the caller used to decide the yield (#2874), threaded through so the logged trigger matches the actual yield reason (a capped-coalesce dispatch logs @@ -478,6 +492,20 @@ def _extract_batch( # Sort pending by cost (ascending) for optimal batching self._pending.sort(key=lambda r: r.item.cost) + reserved_turn = self._reserved_request is not None and bool(self._pending) and not self._pending[0].runs_alone + reserved_limit = len(self._pending) + if reserved_turn: + # Keep flagged-lane ordering intact. The batchable suffix stays + # cost-sorted, and the cheaper prefix waits for the ordinary turn. + reserved_index = next( + index for index, request in enumerate(self._pending) if request is self._reserved_request + ) + batchable_end = next( + (index for index, request in enumerate(self._pending) if request.runs_alone), len(self._pending) + ) + self._pending[:batchable_end] = self._pending[reserved_index:batchable_end] + self._pending[:reserved_index] + reserved_limit = batchable_end - reserved_index + # Greedily take items until we exceed max_batch_cost batch_items: list[I] = [] batch_metadata: list[T] = [] @@ -487,8 +515,10 @@ def _extract_batch( hit_request_limit = False for request in self._pending: + if take_count >= reserved_limit: + break # Respect an explicit per-extract item cap (a drain-fairness bound) - # before anything else, so post-snapshot arrivals stay queued. + # even when new arrivals have joined the pending list. if max_items is not None and take_count >= max_items: break # A runs_alone item never shares a batch: stop in front of one @@ -520,6 +550,16 @@ def _extract_batch( if take_count > 0: # Whose turn is next when both lanes are pending. self._last_batch_alone = lane == "alone" + if lane == "batched": + if reserved_turn: + self._reserved_request = None + elif take_count < len(self._pending): + oldest = min( + (request for request in self._pending if not request.runs_alone), + key=lambda request: request._arrival_order, + ) + if all(request is not oldest for request in self._pending[:take_count]): + self._reserved_request = oldest # Remove taken items from pending self._pending = self._pending[take_count:] @@ -546,6 +586,7 @@ def _extract_batch( if not self._pending: self._first_request_time = None self._last_submit_time = None + self._reserved_request = None return FormattedBatch( items=batch_items, @@ -583,6 +624,7 @@ def drain_pending(self) -> list[T]: self._total_cost = 0 self._first_request_time = None self._last_submit_time = None + self._reserved_request = None return metadata async def try_drain(self, max_items: int | None = None) -> FormattedBatch[I, T] | None: @@ -594,9 +636,10 @@ async def try_drain(self, max_items: int | None = None) -> FormattedBatch[I, T] inference don't need to wait an additional batch formation window. Args: - max_items: Optional cap on how many items this drain may take, so a - caller draining a snapshot backlog does not sweep in items that - arrived mid-drain. See #1606. + max_items: Optional cap on how many items this drain may take from + the current pending list. A caller can bound its total drain + to an earlier backlog count, without freezing member identities. + See #1606. Returns: FormattedBatch if items were pending, None otherwise. 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 066aa54e1..097c10094 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 @@ -1377,7 +1377,7 @@ async def _get_next_batch_fcfs( async def _drain_active_batcher(self, active_lora: str | None) -> bool: """Continuous-batching drain of ``active_lora``'s batcher. - Bounded to the backlog present at entry: a saturated LoRA whose + Bounded to the backlog count at entry: a saturated LoRA whose ``submit()`` coroutines keep refilling its own batcher during each GPU forward pass would otherwise never let this loop exit, starving sibling LoRAs / the base model (they are only re-examined once control returns @@ -1391,10 +1391,9 @@ async def _drain_active_batcher(self, active_lora: str | None) -> bool: if batcher is None: return False - # Snapshot the backlog at entry; items that arrive mid-drain wait for - # the next FCFS selection instead of extending the drain indefinitely. - # Cap every drain to the remaining budget so a batch formed late in the - # loop can't sweep in post-snapshot arrivals and overshoot the backlog. + # Snapshot the backlog count, not its member identities. New arrivals + # may be selected, but cannot extend the total number of items drained + # before the next FCFS selection. Cap every batch to the remaining count. drain_budget = batcher.pending_count drained = False while drain_budget > 0: diff --git a/packages/sie_server/tests/core/test_batcher.py b/packages/sie_server/tests/core/test_batcher.py index 87a418e31..aa35e006f 100644 --- a/packages/sie_server/tests/core/test_batcher.py +++ b/packages/sie_server/tests/core/test_batcher.py @@ -1091,8 +1091,8 @@ class TestRunsAloneLane: """Items flagged ``runs_alone`` (long-form audio, #585) never share a batch.""" @pytest.mark.parametrize("seed", range(20)) - def test_unflagged_items_batch_exactly_as_cost_sorted_greedy_packing(self, seed: int) -> None: - """Without a flagged item pending, every batch matches the old algorithm.""" + def test_ordinary_unflagged_turns_match_cost_sorted_greedy_packing(self, seed: int) -> None: + """Ordinary turns preserve greedy packing; reserved turns serve the oldest.""" rng = random.Random(seed) # noqa: S311 -- reproducible test data config = BatchConfig(max_batch_cost=rng.randint(1, 4_000), max_batch_requests=rng.randint(1, 8)) batcher: BatchFormer[HasCost, int] = BatchFormer(config) @@ -1110,10 +1110,20 @@ def test_unflagged_items_batch_exactly_as_cost_sorted_greedy_packing(self, seed: if not reference: continue max_items = rng.choice([None, None, 1, 2, 5]) + reserved = batcher._reserved_request is not None + oldest = min(request_id for _, request_id in reference) + expected = _greedy_reference(list(reference), config, max_items) batch = batcher._extract_batch(max_items=max_items) - assert batch.metadata == _greedy_reference(reference, config, max_items) + if reserved: + assert oldest in batch.metadata + else: + assert batch.metadata == expected + assert [item.cost for item in batch.items] == sorted(item.cost for item in batch.items) + assert batch.total_cost <= config.max_batch_cost or batch.size == 1 + reference = [entry for entry in reference if entry[1] not in batch.metadata] + reference.sort(key=lambda entry: entry[0]) assert batcher.pending_count == len(reference) assert batcher.pending_cost == sum(cost for cost, _ in reference) @@ -1274,3 +1284,214 @@ def extract() -> None: assert sorted(served) == sorted(submitted) assert not flagged_waits assert batcher.pending_cost == 0 + + +class TestBatchableProgress: + """Replenishing cheaper work cannot indefinitely bypass an older valid item.""" + + @pytest.mark.parametrize(("budget", "heavy_cost"), [(4096, 12_000), (16_384, 12_000), (16_384, 24_000)]) + @pytest.mark.asyncio + async def test_replenishing_cheap_items_cannot_starve_oldest(self, budget: int, heavy_cost: int) -> None: + batcher: BatchFormer[_Audio, str] = BatchFormer(BatchConfig(max_batch_cost=budget)) + await batcher.submit(_Audio(heavy_cost), "heavy") + + served: list[list[str]] = [] + for turn in range(4): + await batcher.submit_many([(_Audio(budget // 64), f"cheap-{turn}-{index}") for index in range(64)]) + batch = await batcher.try_drain() + assert batch is not None + served.append(batch.metadata) + assert batch.total_cost <= budget or batch.size == 1 + + assert "heavy" not in served[0] + assert "heavy" in served[1] + assert all(name.startswith("cheap-") for name in served[2]) + + @pytest.mark.parametrize("cap", [None, 1, 3]) + @pytest.mark.asyncio + async def test_two_heavy_documents_make_bounded_progress(self, cap: int | None) -> None: + """All 48 individually oversized rows progress under continuous cheap arrivals.""" + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(BatchConfig(max_batch_cost=16_384)) + heavy = [ + (ExtractPreparedItem(cost, index), f"long{labels}/{index}") + for labels, cost in ((7, 22_355), (8, 25_554)) + for index in range(24) + ] + await batcher.submit_many(heavy) + served: list[str] = [] + for turn in range(96): + await batcher.submit_many( + [(ExtractPreparedItem(83, index), f"cheap/{turn}/{index}") for index in range(64)] + ) + batch = await batcher.try_drain(max_items=cap) + assert batch is not None + if cap is not None: + assert batch.size <= cap + assert batch.total_cost <= 16_384 or batch.size == 1 + served.extend(batch.metadata) + for position, (_, name) in enumerate(heavy): + if turn + 1 >= 2 * (position + 1): + assert name in served + assert batcher.pending_cost == sum(request.item.cost for request in batcher._pending) + assert all(served.count(name) == 1 for _, name in heavy) + assert any(name.startswith("cheap/") for name in served) + remaining = batcher.drain_pending() + assert len(served + remaining) == len(set(served + remaining)) == 48 + 96 * 64 + + def test_reserved_suffix_packs_peers_without_wrapping_to_cheaper_items(self) -> None: + batcher: BatchFormer[_Audio, str] = BatchFormer(BatchConfig(max_batch_cost=10)) + batcher._append_item(_Audio(4), "old") + batcher._append_item(_Audio(4), "peer") + for index in range(3): + batcher._append_item(_Audio(3), f"cheap-{index}") + assert batcher._extract_batch().metadata == ["cheap-0", "cheap-1", "cheap-2"] + batcher._append_item(_Audio(1), "new-cheap") + + reserved = batcher._extract_batch() + + assert reserved.metadata == ["old", "peer"] + assert reserved.total_cost == 8 + assert batcher._extract_batch().metadata == ["new-cheap"] + + def test_same_timestamp_uses_arrival_identity_not_metadata_equality(self) -> None: + class Metadata: + __hash__ = object.__hash__ + + def __eq__(self, other: object) -> bool: + raise AssertionError("batch selection must never compare metadata") + + batcher: BatchFormer[_Audio, Metadata] = BatchFormer(BatchConfig(max_batch_cost=10)) + first, second = Metadata(), Metadata() + for item, metadata in [(_Audio(20), first), (_Audio(15), second)]: + batcher._append_item(item, metadata) + batcher._pending[-1].arrival_time = 1.0 + + served = [] + for _ in range(4): + batcher._append_item(_Audio(10), Metadata()) + batcher._pending[-1].arrival_time = 1.0 + served.append(batcher._extract_batch().metadata) + assert served[1][0] is first + assert served[3][0] is second + + @pytest.mark.parametrize("new_budget", [5, 12, 100]) + def test_reservation_survives_adaptive_budget_changes(self, new_budget: int) -> None: + config = BatchConfig(max_batch_cost=10) + batcher: BatchFormer[_Audio, str] = BatchFormer(config) + batcher._append_item(_Audio(12), "old") + batcher._append_item(_Audio(10), "cheap") + assert batcher._extract_batch().metadata == ["cheap"] + config.max_batch_cost = new_budget + batcher._append_item(_Audio(1), "new-cheap") + + batch = batcher._extract_batch() + + assert batch.metadata == ["old"] + assert batch.total_cost <= new_budget or batch.size == 1 + assert batcher._extract_batch().metadata == ["new-cheap"] + + @pytest.mark.asyncio + async def test_flagged_and_zero_item_turns_do_not_consume_reservation(self) -> None: + batcher: BatchFormer[_Audio, str] = BatchFormer(BatchConfig(max_batch_cost=10)) + await batcher.submit_many([(_Audio(20), "old"), (_Audio(10), "cheap"), (_Audio(30, runs_alone=True), "alone")]) + first = await batcher.try_drain() + assert first is not None + assert first.metadata == ["cheap"] + reserved = batcher._reserved_request + alone = await batcher.try_drain() + assert alone is not None + assert alone.metadata == ["alone"] + assert batcher._reserved_request is reserved + zero = await batcher.try_drain(max_items=0) + assert zero is not None + assert zero.size == 0 + assert batcher._reserved_request is reserved + await batcher.submit_many([(_Audio(1), "new-cheap"), (_Audio(40, runs_alone=True), "next-alone")]) + + served = [] + while batcher.pending_count: + batch = await batcher.try_drain() + assert batch is not None + served.append(batch.metadata) + assert served == [["old"], ["next-alone"], ["new-cheap"]] + + def test_shutdown_clears_reservation_without_losing_metadata(self) -> None: + batcher: BatchFormer[_Audio, str] = BatchFormer(BatchConfig(max_batch_cost=10)) + batcher._append_item(_Audio(20), "old") + batcher._append_item(_Audio(10), "cheap") + assert batcher._extract_batch().metadata == ["cheap"] + assert batcher._reserved_request is not None + + assert batcher.drain_pending() == ["old"] + assert batcher._reserved_request is None + assert batcher.pending_count == batcher.pending_cost == 0 + batcher._append_item(_Audio(8), "new-large") + batcher._append_item(_Audio(1), "new-small") + assert batcher._extract_batch().metadata == ["new-small", "new-large"] + + @pytest.mark.parametrize("seed", range(10)) + def test_replenishment_respects_arrival_bound_and_conserves_every_item(self, seed: int) -> None: + rng = random.Random(seed) # noqa: S311 -- deterministic adversarial arrivals + batcher: BatchFormer[_Audio, int] = BatchFormer(BatchConfig(max_batch_cost=100, max_batch_requests=8)) + waits: dict[int, list[int]] = {} + served: list[int] = [] + next_id = 0 + + def extract() -> None: + batch = batcher._extract_batch(max_items=rng.choice([None, 0, 1, 3])) + if not batch.size: + return + for wait in waits.values(): + wait[1] += 1 + for name in batch.metadata: + older, turns = waits.pop(name) + assert turns <= 2 * (older + 1) + assert [item.cost for item in batch.items] == sorted(item.cost for item in batch.items) + assert batch.total_cost <= batcher.config.max_batch_cost or batch.size == 1 + served.extend(batch.metadata) + assert batcher.pending_count == len(waits) + assert batcher.pending_cost == sum(request.item.cost for request in batcher._pending) + + for _ in range(100): + for _ in range(rng.randint(1, 5)): + waits[next_id] = [len(waits), 0] + batcher._append_item(_Audio(rng.randint(1, 200)), next_id) + batcher._pending[-1].arrival_time = 1.0 + next_id += 1 + batcher.config.max_batch_cost = rng.choice([50, 100, 300]) + extract() + while batcher.pending_count: + extract() + assert sorted(served) == list(range(next_id)) + + @pytest.mark.parametrize("max_items", [None, 1, 7]) + def test_homogeneous_memberships_and_batch_counts_match_original(self, max_items: int | None) -> None: + config = BatchConfig(max_batch_cost=80, max_batch_requests=16) + batcher: BatchFormer[_Audio, int] = BatchFormer(config) + reference: list[tuple[int, int]] = [] + for turn in range(20): + for index in range(32): + name = turn * 32 + index + batcher._append_item(_Audio(5), name) + reference.append((5, name)) + assert batcher._extract_batch(max_items=max_items).metadata == _greedy_reference( + reference, config, max_items + ) + assert batcher._reserved_request is None + while reference: + assert batcher._extract_batch(max_items=max_items).metadata == _greedy_reference( + reference, config, max_items + ) + + @pytest.mark.parametrize("seed", range(10)) + def test_fitting_mixed_cost_batches_match_original(self, seed: int) -> None: + rng = random.Random(seed) # noqa: S311 -- deterministic fitting batches + config = BatchConfig(max_batch_cost=128, max_batch_requests=16) + batcher: BatchFormer[_Audio, int] = BatchFormer(config) + for _ in range(20): + pending = [(rng.randint(1, 8), index) for index in range(16)] + for cost, name in pending: + batcher._append_item(_Audio(cost), name) + assert batcher._extract_batch().metadata == _greedy_reference(pending, config, None) + assert batcher.pending_count == 0 + assert batcher._reserved_request is None