diff --git a/packages/sie_server/src/sie_server/core/batcher.py b/packages/sie_server/src/sie_server/core/batcher.py index c4a5cf4ca..488b87c11 100644 --- a/packages/sie_server/src/sie_server/core/batcher.py +++ b/packages/sie_server/src/sie_server/core/batcher.py @@ -25,8 +25,27 @@ 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. +When nothing pending is flagged and every pending item fits the configured +``max_batch_cost``, batch formation is exactly the cost-sorted greedy +packing above. + +An item whose cost is above that configured limit cannot share a batch +formed at the limit. Sorted packing never selects it while a cheaper item +is pending, so cheaper arrivals hold it until they stop (#390). Those +items take turns with the cheaper lane: one cost-sorted batch of items +that fit, then an arrival-ordered batch of over-budget items. The first batch +comes from the lane of the oldest pending item. An over-budget item with +k over-budget items ahead of it waits for at most k + 1 cheaper batches, +however many cheaper items keep arriving. The adaptive controller may +later change the live ``max_batch_cost``; lane membership stays on the +budget ``BatchConfig`` was built with, so a raised live limit does not +pull the item back into a cheaper pack. An over-budget turn takes an arrival +prefix that fits the live cost and item limits, or one item if it exceeds the +live cost limit. Items at or under the configured limit stay on cost order, +including when cheaper items fill the live +budget ahead of them. While a ``runs_alone`` item is pending, the existing +flagged/batchable alternation applies and over-budget fairness pauses. The +k + 1 cheaper-batch bound does not apply during that time. Cost semantics vary by modality (modality-native units): - Text: cost = token count @@ -86,13 +105,14 @@ class PendingRequest[I: HasCost, T]: metadata: T arrival_time: float = field(default_factory=time.monotonic) runs_alone: bool = False + arrival_order: int = 0 @dataclass(slots=True) class FormattedBatch[I: HasCost, T]: """A batch ready for inference. - Items are sorted by cost (ascending) to minimize padding. + Items follow the selected lane's cost or arrival order. Metadata is preserved for routing results back to requesters. """ @@ -177,6 +197,7 @@ def __init__( self.max_batch_cost = max_batch_tokens else: self.max_batch_cost = max_batch_cost + self.configured_max_batch_cost = self.max_batch_cost self.max_batch_requests = max_batch_requests self.max_batch_wait_ms = max_batch_wait_ms self.coalesce_ms = coalesce_ms @@ -229,6 +250,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._last_over_budget_batch: bool | None = None + self._next_arrival_order = 0 @property def config(self) -> BatchConfig: @@ -322,7 +345,13 @@ 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 = PendingRequest( + item=item, + metadata=metadata, + runs_alone=getattr(item, "runs_alone", False) is True, + arrival_order=self._next_arrival_order, + ) + self._next_arrival_order += 1 self._pending.append(request) self._total_cost += item.cost @@ -367,8 +396,8 @@ async def get_batch( total wait. Returns: - FormattedBatch with items sorted by cost (ascending). - The batch respects max_batch_cost limit for optimal GPU utilization. + FormattedBatch in the selected lane's cost or arrival order. + The batch respects max_batch_cost, except for a single larger item. """ while True: async with self._lock: @@ -437,6 +466,52 @@ def _order_runs_alone_lane(self) -> None: if alone_turn: self._pending.insert(0, self._pending.pop(head_index)) + def _order_over_budget_lane(self) -> bool: + """Order lanes and return whether this turn serves over-budget items. + + Packable items stay cost-sorted. Over-budget items stay in arrival + order and are not taken during a packable turn, so a raised live + ``max_batch_cost`` cannot fold one into the cheaper batch. On an + over-budget turn the whole arrival-ordered lane moves to the front; + the caller packs its prefix under the live cost and item limits. + """ + budget = self._config.configured_max_batch_cost + has_over = False + has_pack = False + for request in self._pending: + if request.item.cost > budget: + has_over = True + else: + has_pack = True + if not has_over: + self._pending.sort(key=lambda request: request.item.cost) + return False + + self._pending.sort( + key=lambda request: ( + request.item.cost > budget, + request.arrival_time if request.item.cost > budget else request.item.cost, + request.arrival_order if request.item.cost > budget else 0, + ) + ) + if has_pack and not self._over_budget_turn(budget): + return False + head_index = next(index for index, request in enumerate(self._pending) if request.item.cost > budget) + if head_index != 0: + self._pending[:] = self._pending[head_index:] + self._pending[:head_index] + return True + + def _over_budget_turn(self, budget: int) -> bool: + """Whether the next batch belongs to the over-budget lane. + + Both lanes are pending. The first batch follows the oldest item. + After that the lanes alternate. + """ + if self._last_over_budget_batch is None: + oldest = min(self._pending, key=lambda request: (request.arrival_time, request.arrival_order)) + return oldest.item.cost > budget + return not self._last_over_budget_batch + def _extract_batch( self, max_items: int | None = None, @@ -445,7 +520,7 @@ def _extract_batch( """Extract an optimal sub-batch from pending requests. Algorithm: - 1. Sort pending requests by cost (ascending) + 1. Order pending requests by the selected lane's cost or arrival order 2. Take requests until max_batch_cost would be exceeded 3. Keep remaining requests for next get_batch() call @@ -453,6 +528,8 @@ def _extract_batch( 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. + Otherwise, items above the configured cost limit alternate with + cost-sorted batches, packing in arrival order under the live limits. Args: max_items: Optional hard cap on the number of items taken, on top of @@ -472,11 +549,12 @@ def _extract_batch( pending_before = len(self._pending) cost_before = self._total_cost - if any(request.runs_alone for request in self._pending): + alone_lane = any(request.runs_alone for request in self._pending) + over_budget_turn = False + if alone_lane: self._order_runs_alone_lane() else: - # Sort pending by cost (ascending) for optimal batching - self._pending.sort(key=lambda r: r.item.cost) + over_budget_turn = self._order_over_budget_lane() # Greedily take items until we exceed max_batch_cost batch_items: list[I] = [] @@ -495,6 +573,14 @@ def _extract_batch( # unless it is the first item taken. if request.runs_alone and take_count > 0: break + # Keep the over-budget and cheaper lanes separate. Flagged turns + # retain the existing runs_alone scheduler instead. + if ( + not alone_lane + and take_count > 0 + and (request.item.cost > self._config.configured_max_batch_cost) != over_budget_turn + ): + break # Would this item push us over the cost limit? # Always take at least one item (handles single large requests) if batch_cost + request.item.cost > self._config.max_batch_cost: @@ -516,10 +602,17 @@ def _extract_batch( hit_request_limit = True break - lane = "alone" if take_count > 0 and self._pending[0].runs_alone else "batched" + if take_count > 0 and self._pending[0].runs_alone: + lane = "alone" + elif over_budget_turn: + lane = "over_budget" + else: + lane = "batched" if take_count > 0: # Whose turn is next when both lanes are pending. self._last_batch_alone = lane == "alone" + if not alone_lane: + self._last_over_budget_batch = over_budget_turn # Remove taken items from pending self._pending = self._pending[take_count:] @@ -542,10 +635,12 @@ def _extract_batch( len(self._pending), ) - # Reset timers if we've emptied the queue + # Reset timers if we've emptied the queue. The next queue follows its + # oldest item. The runs_alone lane keeps its turn across an empty queue. if not self._pending: self._first_request_time = None self._last_submit_time = None + self._last_over_budget_batch = None return FormattedBatch( items=batch_items, @@ -557,7 +652,7 @@ def try_get_batch(self) -> FormattedBatch[I, T] | None: """Try to get a batch without waiting. Returns: - FormattedBatch if batch is ready (sorted by cost), None otherwise. + FormattedBatch in the selected lane's order if ready, None otherwise. """ if self._should_yield_batch(): return self._extract_batch() @@ -583,6 +678,7 @@ def drain_pending(self) -> list[T]: self._total_cost = 0 self._first_request_time = None self._last_submit_time = None + self._last_over_budget_batch = None return metadata async def try_drain(self, max_items: int | None = None) -> FormattedBatch[I, T] | None: diff --git a/packages/sie_server/tests/core/test_batcher.py b/packages/sie_server/tests/core/test_batcher.py index 87a418e31..e0910f55b 100644 --- a/packages/sie_server/tests/core/test_batcher.py +++ b/packages/sie_server/tests/core/test_batcher.py @@ -1092,16 +1092,21 @@ class TestRunsAloneLane: @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.""" + """Items that fit the configured budget match cost-sorted greedy packing. + + An item above that budget uses the over-budget lane instead. This + control keeps every generated cost inside the budget. + """ 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)) + configured_budget = config.max_batch_cost batcher: BatchFormer[HasCost, int] = BatchFormer(config) reference: list[tuple[int, int]] = [] next_id = 0 for _ in range(60): for _ in range(rng.randint(0, 5)): - cost = rng.randint(1, 1_500) + cost = rng.randint(1, configured_budget) # Items with runs_alone=False and items without the attribute. item: HasCost = _Audio(cost=cost) if rng.random() < 0.5 else ExtractPreparedItem(cost, 0) batcher._append_item(item, next_id) @@ -1109,7 +1114,9 @@ def test_unflagged_items_batch_exactly_as_cost_sorted_greedy_packing(self, seed: next_id += 1 if not reference: continue - max_items = rng.choice([None, None, 1, 2, 5]) + config.max_batch_cost = rng.randint(1, configured_budget * 2) + config.max_batch_requests = rng.randint(1, 8) + max_items = rng.choice([None, 0, 1, 2, 5]) batch = batcher._extract_batch(max_items=max_items) @@ -1274,3 +1281,362 @@ def extract() -> None: assert sorted(served) == sorted(submitted) assert not flagged_waits assert batcher.pending_cost == 0 + + +def _queue(batcher: BatchFormer[ExtractPreparedItem, str], cost: int, name: str, *, at: float) -> None: + """Queue an extract row as if it arrived at monotonic time ``at``.""" + batcher._append_item(ExtractPreparedItem(cost, 0), name) + batcher._pending[-1].arrival_time = at + + +class TestOverBudgetLane: + """Items above the configured batch cost take turns with cheaper items (#390). + + Cost-sorted packing never reaches an item that does not fit the configured + limit while a cheaper item is pending. The cheaper lane keeps one batch, + then an arrival-ordered over-budget prefix fits the live limits. Items + that fit stay on cost order. Waiting is not itself a reason to move an item. + """ + + def test_over_budget_items_dispatch_while_cheaper_items_keep_arriving(self) -> None: + """24 rows above the GLiClass budget move while cheaper rows are still queued.""" + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer( + BatchConfig(max_batch_cost=16_384, max_batch_requests=64) + ) + for index in range(24): + _queue(batcher, 22_355 + index, f"long-{index}", at=float(index)) + for index in range(8): + _queue(batcher, 128, f"cheap-{index}", at=100.0 + index) + + served: list[list[str]] = [] + for index in range(48): + served.append(batcher._extract_batch().metadata) + _queue(batcher, 128, f"later-{index}", at=200.0 + index) + + long_batches = [batch for batch in served if batch[0].startswith("long-")] + assert [batch[0] for batch in long_batches] == [f"long-{index}" for index in range(24)] + assert all(len(batch) == 1 for batch in long_batches) + cheap_batches = [batch for batch in served if not batch[0].startswith("long-")] + assert all(not name.startswith("long-") for batch in cheap_batches for name in batch) + cheap_before = 0 + seen_long = 0 + for batch in served: + if batch[0].startswith("long-"): + assert cheap_before <= seen_long + 1 + seen_long += 1 + else: + cheap_before += 1 + assert batcher.pending_count > 0 + + def test_the_oldest_pending_item_decides_the_first_batch(self) -> None: + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(BatchConfig(max_batch_cost=100)) + _queue(batcher, 250, "long", at=0.0) + _queue(batcher, 40, "cheap", at=1.0) + + assert batcher._extract_batch().metadata == ["long"] + assert batcher._extract_batch().metadata == ["cheap"] + + def test_an_empty_queue_restarts_the_over_budget_lane(self) -> None: + """The next queue follows its oldest item after an extract empties the queue.""" + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(BatchConfig(max_batch_cost=100)) + _queue(batcher, 250, "lone", at=0.0) + assert batcher._extract_batch().metadata == ["lone"] + assert batcher.pending_count == 0 + + _queue(batcher, 250, "over-budget", at=0.0) + _queue(batcher, 40, "cheap", at=1.0) + + assert batcher._extract_batch().metadata == ["over-budget"] + assert batcher._extract_batch().metadata == ["cheap"] + + def test_drain_pending_restarts_the_over_budget_lane(self) -> None: + """Draining the remainder of a queue clears the over-budget lane flag.""" + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(BatchConfig(max_batch_cost=100)) + _queue(batcher, 250, "over-budget", at=0.0) + _queue(batcher, 40, "cheap", at=1.0) + assert batcher._extract_batch().metadata == ["over-budget"] + + assert batcher.drain_pending() == ["cheap"] + assert batcher.pending_count == 0 + + _queue(batcher, 250, "next-over-budget", at=0.0) + _queue(batcher, 40, "next-cheap", at=1.0) + assert batcher._extract_batch().metadata == ["next-over-budget"] + + def test_over_budget_items_run_in_arrival_order(self) -> None: + """A newer, cheaper over-budget row does not pass an older one.""" + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(BatchConfig(max_batch_cost=100)) + _queue(batcher, 400, "older-larger", at=0.0) + _queue(batcher, 200, "newer-smaller", at=1.0) + + assert batcher._extract_batch().metadata == ["older-larger"] + assert batcher._extract_batch().metadata == ["newer-smaller"] + + def test_items_within_the_configured_budget_stay_on_cost_order(self) -> None: + """A row that fits is not pulled ahead because cheaper rows filled the budget.""" + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer( + BatchConfig(max_batch_cost=100, max_batch_requests=10) + ) + _queue(batcher, 80, "larger", at=0.0) + for index in range(6): + _queue(batcher, 30, f"cheap-{index}", at=float(index + 1)) + + assert batcher._extract_batch().metadata == ["cheap-0", "cheap-1", "cheap-2"] + for index in range(3): + _queue(batcher, 30, f"more-{index}", at=10.0 + index) + second = batcher._extract_batch().metadata + assert "larger" not in second + assert all(name != "larger" for name in second) + + def test_a_raised_live_budget_does_not_pack_the_over_budget_item_with_cheaper_ones(self) -> None: + config = BatchConfig(max_batch_cost=16_384) + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(config) + _queue(batcher, 100, "cheap", at=0.0) + _queue(batcher, 22_355, "long", at=1.0) + config.max_batch_cost = 65_536 + + assert config.configured_max_batch_cost == 16_384 + assert batcher._extract_batch().metadata == ["cheap"] + assert batcher._extract_batch().metadata == ["long"] + assert batcher.pending_cost == 0 + + @pytest.mark.parametrize("with_cheap", [False, True]) + def test_raised_budget_packs_an_over_budget_fifo_prefix(self, with_cheap: bool) -> None: + config = BatchConfig(max_batch_cost=100) + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(config) + for index, cost in enumerate([200, 150, 110]): + _queue(batcher, cost, f"long-{index}", at=float(index)) + if with_cheap: + _queue(batcher, 1, "cheap", at=3.0) + config.max_batch_cost = 450 + + first = batcher._extract_batch() + + assert first.metadata == ["long-0", "long-1"] + assert first.total_cost == 350 + if with_cheap: + assert batcher._extract_batch().metadata == ["cheap"] + assert batcher._extract_batch().metadata == ["long-2"] + assert batcher.pending_count == batcher.pending_cost == 0 + + def test_fifo_packing_does_not_skip_an_older_item_to_fill_spare_budget(self) -> None: + config = BatchConfig(max_batch_cost=100) + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(config) + for index, cost in enumerate([140, 190, 110]): + _queue(batcher, cost, f"long-{index}", at=float(index)) + config.max_batch_cost = 300 + + assert batcher._extract_batch().metadata == ["long-0"] + assert batcher._extract_batch().metadata == ["long-1", "long-2"] + + @pytest.mark.parametrize(("request_cap", "item_cap", "taken"), [(1, None, 1), (8, 2, 2), (3, 8, 3)]) + def test_over_budget_fifo_pack_respects_both_item_caps(self, request_cap, item_cap, taken) -> None: + config = BatchConfig(max_batch_cost=100, max_batch_requests=request_cap) + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(config) + for index in range(4): + _queue(batcher, 200 - index, f"long-{index}", at=float(index)) + config.max_batch_cost = 1_000 + + batch = batcher._extract_batch(max_items=item_cap) + + assert batch.metadata == [f"long-{index}" for index in range(taken)] + assert batcher.pending_count == 4 - taken + assert batcher.pending_cost + batch.total_cost == 794 + + def test_lowered_live_budget_still_dispatches_one_oversized_item(self) -> None: + config = BatchConfig(max_batch_cost=100) + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(config) + _queue(batcher, 250, "first", at=0.0) + _queue(batcher, 150, "second", at=1.0) + config.max_batch_cost = 50 + + assert batcher._extract_batch().metadata == ["first"] + assert batcher._extract_batch().metadata == ["second"] + assert batcher.pending_count == batcher.pending_cost == 0 + + @pytest.mark.parametrize("over_first", [False, True]) + def test_tied_timestamps_choose_the_first_submitted_lane(self, over_first: bool) -> None: + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(BatchConfig(max_batch_cost=100)) + rows = [(250, "long"), (40, "cheap")] + if not over_first: + rows.reverse() + for cost, name in rows: + _queue(batcher, cost, name, at=1.0) + + assert batcher._extract_batch().metadata == [rows[0][1]] + assert batcher._extract_batch().metadata == [rows[1][1]] + + @pytest.mark.parametrize("live_budget", [100, 600]) + def test_tied_over_budget_arrivals_survive_runs_alone_cost_sorting(self, live_budget: int) -> None: + config = BatchConfig(max_batch_cost=100) + batcher: BatchFormer[_Audio, str] = BatchFormer(config) + _arrive(batcher, _Audio(1_000, runs_alone=True), "alone", at=0.0) + _arrive(batcher, _Audio(300), "first", at=1.0) + _arrive(batcher, _Audio(200), "second", at=1.0) + _arrive(batcher, _Audio(10), "cheap", at=2.0) + config.max_batch_cost = live_budget + + assert batcher._extract_batch().metadata == ["alone"] + if live_budget == 100: + assert batcher._extract_batch().metadata == ["first"] + assert batcher._extract_batch().metadata == ["cheap"] + assert batcher._extract_batch().metadata == ["second"] + else: + assert batcher._extract_batch().metadata == ["first", "second"] + assert batcher._extract_batch().metadata == ["cheap"] + assert batcher.pending_count == batcher.pending_cost == 0 + + def test_flagged_lane_preserves_normal_live_budget_packing(self) -> None: + config = BatchConfig(max_batch_cost=100) + batcher: BatchFormer[_Audio, str] = BatchFormer(config) + _arrive(batcher, _Audio(10), "cheap", at=0.0) + _arrive(batcher, _Audio(200), "large", at=1.0) + _arrive(batcher, _Audio(1_000, runs_alone=True), "alone", at=2.0) + config.max_batch_cost = 500 + + assert batcher._extract_batch().metadata == ["cheap", "large"] + assert batcher._extract_batch().metadata == ["alone"] + + def test_zero_and_flagged_turns_preserve_the_next_over_budget_turn(self) -> None: + config = BatchConfig(max_batch_cost=100) + batcher: BatchFormer[_Audio, str] = BatchFormer(config) + _arrive(batcher, _Audio(40), "cheap", at=0.0) + _arrive(batcher, _Audio(200), "first", at=1.0) + _arrive(batcher, _Audio(150), "second", at=2.0) + assert batcher._extract_batch().metadata == ["cheap"] + _arrive(batcher, _Audio(1_000, runs_alone=True), "alone", at=3.0) + _arrive(batcher, _Audio(1), "later-cheap", at=4.0) + config.max_batch_cost = 500 + + assert batcher._extract_batch().metadata == ["alone"] + assert batcher._extract_batch(max_items=0).metadata == [] + assert batcher._extract_batch().metadata == ["first", "second"] + assert batcher._extract_batch().metadata == ["later-cheap"] + + def test_over_budget_order_never_compares_metadata(self) -> None: + class Metadata: + __hash__ = object.__hash__ + + def __eq__(self, other: object) -> bool: + raise AssertionError("batch selection must not compare metadata") + + batcher: BatchFormer[ExtractPreparedItem, Metadata] = BatchFormer(BatchConfig(max_batch_cost=100)) + first, second = Metadata(), Metadata() + for cost, metadata in [(250, first), (150, second)]: + batcher._append_item(ExtractPreparedItem(cost, 0), metadata) + batcher._pending[-1].arrival_time = 1.0 + batcher.config.max_batch_cost = 500 + + batch = batcher._extract_batch() + + assert len(batch.metadata) == 2 + assert batch.metadata[0] is first + assert batch.metadata[1] is second + + def test_turn_taking_keeps_pending_cost_consistent(self) -> None: + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(BatchConfig(max_batch_cost=100)) + _queue(batcher, 40, "a", at=0.0) + _queue(batcher, 250, "b", at=1.0) + _queue(batcher, 40, "c", at=2.0) + + assert batcher.pending_cost == 330 + first = batcher._extract_batch() + assert first.metadata == ["a", "c"] + assert first.total_cost == 80 + assert batcher.pending_cost == 250 + second = batcher._extract_batch() + assert second.metadata == ["b"] + assert second.total_cost == 250 + assert batcher.pending_count == 0 + assert batcher.pending_cost == 0 + + @pytest.mark.parametrize("seed", range(20)) + def test_over_budget_items_wait_for_at_most_one_extra_cheaper_batch(self, seed: int) -> None: + rng = random.Random(seed) # noqa: S311 -- reproducible test data + budget = rng.randint(10, 500) + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer( + BatchConfig(max_batch_cost=budget, max_batch_requests=rng.randint(1, 8)) + ) + served: list[str] = [] + submitted: list[str] = [] + waits: dict[str, list[int]] = {} + last_over: bool | None = None + + def extract() -> None: + nonlocal last_over + pending = list(batcher._pending) + batch = batcher._extract_batch(max_items=rng.choice([None, 1, 3])) + assert batch.size > 0 + over = batch.items[0].cost > budget + assert batch.total_cost <= batcher.config.max_batch_cost or batch.size == 1 + assert all((item.cost > budget) is over for item in batch.items) + has_over = any(request.item.cost > budget for request in pending) + has_pack = any(request.item.cost <= budget for request in pending) + if has_over and has_pack: + pending_by_name = {request.metadata: request for request in pending} + oldest_name = next(name for name in submitted if name in pending_by_name) + oldest_over = pending_by_name[oldest_name].item.cost > budget + assert over is (oldest_over if last_over is None else not last_over) + last_over = over + if batcher.pending_count == 0: + last_over = None + if over: + assert batch.metadata == list(waits)[: batch.size] + for name in batch.metadata: + ahead, cheap_batches = waits.pop(name) + assert cheap_batches <= ahead + 1 + else: + for wait in waits.values(): + wait[1] += 1 + served.extend(batch.metadata) + assert batch.total_cost == sum(item.cost for item in batch.items) + assert batcher.pending_cost == sum(request.item.cost for request in batcher._pending) + + now = 0.0 + for index in range(200): + now += rng.choice([0.0, rng.uniform(0.0, 2.0)]) + if rng.random() < 0.7: + name = f"item-{index}" + if rng.random() < 0.25: + waits[name] = [len(waits), 0] + cost = budget + rng.randint(1, budget) + else: + cost = rng.randint(1, budget) + _queue(batcher, cost, name, at=now) + submitted.append(name) + if batcher.pending_count and rng.random() < 0.55: + batcher.config.max_batch_cost = rng.choice([max(1, budget // 2), budget, budget * 3]) + extract() + while batcher.pending_count: + extract() + + assert sorted(served) == sorted(submitted) + assert not waits + assert batcher.pending_cost == 0 + + @pytest.mark.asyncio + async def test_try_drain_takes_the_over_budget_turn(self) -> None: + batcher: BatchFormer[ExtractPreparedItem, str] = BatchFormer(BatchConfig(max_batch_cost=100)) + await batcher.submit_many( + [ + (ExtractPreparedItem(40, 0), "cheap-a"), + (ExtractPreparedItem(250, 1), "long-a"), + (ExtractPreparedItem(40, 2), "cheap-b"), + (ExtractPreparedItem(180, 3), "long-b"), + ] + ) + # submit_many stamps arrival in call order. Equal timestamps keep that order. + for index, request in enumerate(batcher._pending): + request.arrival_time = float(index) + + first = await batcher.try_drain() + assert first is not None + assert first.metadata == ["cheap-a", "cheap-b"] + second = await batcher.try_drain() + assert second is not None + assert second.metadata == ["long-a"] + third = await batcher.try_drain(max_items=1) + assert third is not None + assert third.metadata == ["long-b"] + assert await batcher.try_drain() is None + assert batcher.pending_cost == 0 diff --git a/packages/sie_server/tests/core/test_progress_batcher.py b/packages/sie_server/tests/core/test_progress_batcher.py index 0ce381406..83397d2d0 100644 --- a/packages/sie_server/tests/core/test_progress_batcher.py +++ b/packages/sie_server/tests/core/test_progress_batcher.py @@ -150,6 +150,17 @@ def test_reservation_survives_adaptive_budget_changes(self, new_budget: int) -> assert batch.total_cost <= new_budget or batch.size == 1 assert batcher._extract_batch().metadata == ["new-cheap"] + def test_progress_policy_does_not_apply_the_default_over_budget_lane(self) -> None: + config = BatchConfig(max_batch_cost=100) + batcher: ProgressBatchFormer[_Audio, str] = ProgressBatchFormer(config) + batcher._append_item(_Audio(200), "older-large") + batcher._append_item(_Audio(10), "newer-cheap") + config.max_batch_cost = 500 + + assert batcher._extract_batch().metadata == ["newer-cheap", "older-large"] + assert batcher._reserved_request is None + assert batcher._arrivals == {} + @pytest.mark.asyncio async def test_flagged_and_zero_item_turns_do_not_consume_reservation(self) -> None: batcher: ProgressBatchFormer[_Audio, str] = ProgressBatchFormer(BatchConfig(max_batch_cost=10)) diff --git a/packages/sie_server_sidecar/src/scheduler/batch_former.rs b/packages/sie_server_sidecar/src/scheduler/batch_former.rs index ecb451d3e..7b507bdfa 100644 --- a/packages/sie_server_sidecar/src/scheduler/batch_former.rs +++ b/packages/sie_server_sidecar/src/scheduler/batch_former.rs @@ -82,9 +82,11 @@ //! extract and clears them only when empty. Rust re-anchors them to the actual //! remainder because the sidecar's cross-key FCFS selector reads that state: //! retaining a removed request's timestamp can repeatedly prioritize a hot key -//! and can make a newer tail skip its own coalescing window. Python's packing -//! remains strictly cost-sorted; Rust promotes at most one aged request per -//! extract, outside batchable turns (points 3 and 5). +//! and can make a newer tail skip its own coalescing window. Python's default +//! packing alternates cost-sorted fitting items with an arrival-ordered +//! over-budget lane when nothing is flagged. Its optional bounded-progress +//! policy can reserve a bypassed oldest item instead. Rust promotes at most +//! one aged request per extract, outside batchable turns (points 3 and 5). //! //! The runs-alone lane (point 5) is the Python batcher's rule too (#601), //! with three differences: