[kozistr]

Context

• In DynamicBatchScheduler::Enqueue, a full-queue rejection leaves queued_batch_size_ inflated because the scheduler increments it before admission succeeds. Later admissions can trigger unnecessary wake-ups against the preferred batch-size threshold.
• The batcher also reads payload readiness without the execution mutex while dequeue changes that state. The accounting test exposed this race under ThreadSanitizer.
• In RateLimiter, live instance changes erase specific_queues_ while other threads read capacity or look up queues. ThreadSanitizer reproduced map reads racing with node deallocation. Removal can destroy a queue that a blocking slot check still uses.
• The scheduler's slot wait ignores shutdown, so teardown can hang after an accounting-test failure releases the consumer.

Honestly, i'm very new here and not very familar with C++, so i used Codex to help me understand and reproduce the bug. So, please fell free to let me know if there are any issues with the coding style or changes!

Solution

• Capture the request batch size before ownership transfer, then update queued_batch_size_ after successful admission.
• Extend the existing execution-mutex block to cover the payload readiness check. Keep enqueue outside the lock.
• Use PayloadQueue::mu_ for instance-map access and removal. Move consumer-count updates into the existing enqueue/dequeue critical sections, preserving accounting for targeted and generic payloads. Run payload-release and final dequeue callbacks outside that mutex.
• Retain queue ownership through blocking waits. Close queues during removal, wake their waiters, and return false from cancelled slot checks. Release the registry and payload mutexes before waiting.
• Let the scheduler leave its slot wait during shutdown. On accounting-test exceptions, wake and join the active consumer before rethrowing the original failure.

Other things

I compared median throughput across three runs per build, with 10,000 measured requests per case and 32 HTTP clients:

CPU Python-backend case	Baseline req/s	Fixed req/s	Change
Unlimited queue, max batch size 1	5776.4	5855.0	+1.36%
Finite queue, max batch size 8	5797.2	5806.0	+0.15%

Throughput differences are smaller than the variation across runs. These measurements cover CPU HTTP/Python workloads; I have not tested GPU workloads.

Reproduction and environment

Compare base commit 3af839d613da6995051f9bcfab00efe2eac87d4c with this branch. I used nvcr.io/nvidia/tritonserver:26.09-py3 on Docker Desktop/WSL2, Linux x86_64, CPU, GCC 13.3, CMake 3.31.10, and GCC ThreadSanitizer. I built both libraries with the same optimized settings and dependencies, then loaded the fixed library into Triton via LD_LIBRARY_PATH. The local build required the Protobuf compatibility alias -Dalways_print_fields_with_no_presence=always_print_primitive_fields.

Build and run the committed regression harness:

cmake --build BUILD --target scheduler_regression_test
ctest --test-dir BUILD/test -R SchedulerRegression --output-on-failure

The accounting test fills a two-request queue with batches of eight and one, rejects seven admissions, drains the batch of eight, and admits another request. A batching callback checks that the later admission leaves the batcher waiting below the preferred size of eight. Restoring the old increment triggers an extra wake-up and fails the assertion. The test drains both remaining requests before scheduler teardown, including after that assertion fails.

For the map races, configure a separate build with -DCMAKE_CXX_FLAGS="-fsanitize=thread -fno-omit-frame-pointer" and -DCMAKE_EXE_LINKER_FLAGS="-fsanitize=thread -no-pie", then run the same commands. The capacity and lookup cases read queues while another thread registers and removes two instances for 2,000 iterations. Removing the erase lock triggers sanitizer reports. The instance/model-removal cases block a slot check and require removal to release it with a false result.

I instrumented the harness and its core sources. Container address randomization caused some sanitizer startup failures; I retried those starts and treated sanitizer reports or test failures as failures.

[greptile-apps]

Confidence Score: 5/5

[High impact] The PR appears safe to merge; no new actionable issues were found.

Summary

Fixes rejected-request accounting and protects instance queues during concurrent lookups and removal.

• The batcher counts requests only after the queue accepts them.
• Rate-limiter waiters keep their queues and wake when those queues close.

No new actionable issues were found. The code and build setup were reviewed without running the regression tests.

Diagram

sequenceDiagram
  participant Waiter as Slot checker
  participant Limiter as Rate limiter
  participant Queue as Instance queue
  participant Removal as Instance removal
  Waiter->>Limiter: Check for a consumer
  Limiter->>Limiter: Copy shared queue ownership under locks
  Limiter->>Queue: Wait after releasing registry locks
  Removal->>Queue: Close queue
  Queue-->>Waiter: Wake and return false
  Removal->>Limiter: Erase queue entry

Loading

_{Reviews (7) · Last reviewed commit: "Merge branch 'main' into fix/batcher-acc..." · Reviewed by Greptile}

[whoisj]

CI-Pipeline-ID: 75939693