Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions src/aiperf/credit/sticky_router.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
- StickyCreditRouter: Main router class
"""

import asyncio
import time
from collections import defaultdict
from collections.abc import Awaitable, Callable
Expand Down Expand Up @@ -274,6 +275,36 @@ def set_first_token_callback(
"""Set callback for first token events (enables prefill concurrency release)."""
self._on_first_token_callback = callback

async def _await_workers_registered(
self, timeout_s: float = 60.0, poll_s: float = 0.05
) -> None:
"""Wait until at least one worker has registered.

A phase-start burst dispatch can outrun worker registration: workers
announce themselves with ``WorkerReady`` over ZMQ *after* the timing
manager begins issuing credits (measured 0.4-0.9s of skew, widening
with concurrency because more workers must register).

Raising from ``send_credit`` in that window is unrecoverable: the
caller (``AgenticReplayTiming._execute_warmup``) issues credits in a
*sequential* ``await`` loop, so the exception aborts the whole burst and
the remaining credits are never attempted. Worse, ``increment_sent``
has already run for the failed credit, so it stays counted as in-flight
with no timeout -- the run hangs forever with no error surfaced
(server side sees no request at all).

Waiting turns that race into a sub-second delay.
"""
waited = 0.0
while not self._workers and waited < timeout_s:
await asyncio.sleep(poll_s)
waited += poll_s
if self._workers:
self.info(
f"Waited {waited:.2f}s for worker registration "
f"({len(self._workers)} workers) before routing credits"
)

async def send_credit(self, credit: Credit) -> None:
"""Determine the worker based on sticky sessions or least-loaded and send the credit to the worker.

Expand All @@ -282,6 +313,10 @@ async def send_credit(self, credit: Credit) -> None:
- Updates the worker load and sticky sessions
- Sends the credit to the worker
"""
if not self._workers:
# Phase-start bursts can precede worker registration; wait instead
# of raising (see _await_workers_registered for why raising hangs).
await self._await_workers_registered()
if not self._workers:
raise RuntimeError("No workers available for routing")

Expand Down
Loading