fix: prevent dynamic-batcher starvation from waiting_consumer_count drift

[jimmystique]

Relates to triton-inference-server/server#8870

Summary

InstanceQueue::Dequeue can merge multiple queued payloads into one payload. Each payload was already counted as claiming a waiting consumer when it was enqueued, but RateLimiter::DequeuePayload only increments the waiting-consumer count once for the single dequeue call. This leaves a permanent -(k-1) drift whenever k extra payloads are merged.

With prefetching disabled (default_queue_policy.default_timeout_microseconds != 0), the dynamic batcher gates dispatch on WaitingConsumerCount() > 0, so the drift can progressively degrade dispatch and eventually block it until model reload.

This change credits back one waiting-consumer count per merged payload inside InstanceQueue::Dequeue, where the merge happens, so the credit lands on the same queue that was dequeued from.

Why This Location

Doing the rebalance inside InstanceQueue::Dequeue keeps the accounting local to the merge and works for both paths:

• per-model batching: Dequeue runs on queue_, whose count gates dispatch;
• device-blocking batching: Dequeue runs on specific_queues_[instance], whose count gates dispatch.

Testing

Added src/test/instance_queue_test.cc, which reproduces the accounting drift across repeated merge-heavy dequeues and verifies that WaitingConsumerCount() returns to the idle-instance count. Without the fix, the test ends with a negative count.

[Vinya567]

Reviewed and validated.

Root cause looks correct: enqueue decrements waiting_consumer_count_ once per payload, dequeue increments once, so merges can leak credits and starve the dynamic batcher. The fix in InstanceQueue::Dequeue (credit one per merged payload) is in the right place.

Unit test fails without the fix and passes with it (verified locally). Related dynamic-batching / infer paths also look good from my side.

Follow-up: the new test is built/installed in core, but isn’t wired into an L0 runner yet , I’ll open a small follow-up so CI actually runs it after merge.

LGTM from my side.

[yinggeh]

Verdict

Approve. Root cause and fix look correct for the reported starvation path (server#8870).

Why the fix is right

RateLimiter::EnqueuePayload decrements waiting_consumer_count_ once per payload; DequeuePayload increments once per dequeue call. InstanceQueue::Dequeue can merge additional queued payloads into that single dequeue, so each merge of k payloads permanently leaks -(k-1).

With prefetching disabled (default_queue_policy.default_timeout_microseconds != 0), PayloadSlotAvailable / WaitForPayloadSlotAvailable gate on WaitingConsumerCount() > 0, so the drift progressively stalls dispatch until model reload.

Crediting one count per entry in merged_payloads inside InstanceQueue::Dequeue is the right place: the merge is local, and the credit lands on the same queue that was dequeued from (per-model queue_ or device-blocking specific_queues_[instance]), which is exactly what gates that batcher mode.

Locking looks fine: waiting_consumer_mu_ is never held across a payload_queue->mu_ acquisition, so doing IncrementConsumerCount() under the payload-queue lock during Dequeue does not invert the existing order.

Test

ConsumerCountStableAcrossMerges models the RateLimiter contract correctly (idle +N, per-enqueue -1, post-work re-park +1, with merge credits supplied by Dequeue). Empty payloads + non-zero max_queue_delay_ns reliably force merges. The no-merge case is a good guard against over-crediting. Stubbing Payload to avoid pulling scheduler/backend symbols matches the existing lightweight unit-test style (repo_agent_test, etc.).

Noted that L0 wiring is handled separately in server#8905 — good follow-up so CI actually runs the binary.

Nits / follow-ups (non-blocking)

1. Copyright year on modified src/instance_queue.cc is still 2021-2022; please bump to include 2026 before merge (Triton copyright check).

2. Residual queue_ drift on instance-bound payloads (optional follow-up): EnqueuePayload always decrements queue_ and, when payload->GetInstance() != nullptr, also decrements the specific queue. This PR only credits the queue Dequeue runs on. For per-instance batchers the dispatch gate reads the specific count, so the reported bug is fixed; queue_ can still drift under that path. Not needed for #8870, but worth a follow-up if we want the dual-queue counters fully balanced (e.g. also credit queue_ from RateLimiter::DequeuePayload when merges come from a specific queue).

[yinggeh]

Please fix pre-commit pre-commit / pre-commit (pull_request)

[greptile-apps]

Greptile Summary

This PR restores waiting-consumer accounting when InstanceQueue::Dequeue merges payloads.

• Credits the dequeued queue once for each merged payload.
• Adds standalone tests covering repeated merge-heavy dequeues and non-merge behavior.
• Registers the new instance_queue_test target in CMake.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
src/instance_queue.cc	Restores one waiting-consumer credit for every payload successfully merged and removed during a dequeue.
src/test/instance_queue_test.cc	Adds focused regression coverage for stable consumer accounting across merged and unmerged dequeues.
src/test/CMakeLists.txt	Adds the standalone InstanceQueue regression-test executable and its required dependencies.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Enqueue each payload] --> B[Decrement consumer count per payload]
  B --> C[Dequeue primary payload]
  C --> D[Merge additional payloads]
  D --> E[Caller restores one dequeue credit]
  D --> F[InstanceQueue restores one credit per merged payload]
  E --> G[Consumer count returns to idle-instance count]
  F --> G

Loading

_{Reviews (3): Last reviewed commit: "style: apply clang-format" | Re-trigger Greptile}

[Vinya567]

Updated copyright years (src/instance_queue.cc → 2021-2026, src/test/CMakeLists.txt → 2019-2026).