diff --git a/docs/load_generator/DESIGN.md b/docs/load_generator/DESIGN.md index ac9327a7b..a3c999170 100644 --- a/docs/load_generator/DESIGN.md +++ b/docs/load_generator/DESIGN.md @@ -309,9 +309,12 @@ async def _receive_responses(self): if self._strategy_task and not self._strategy_task.done(): self._strategy_task.cancel() break - self._handle_response(resp) + # One clock read per response, shared by the chunk event records + # (RECV_FIRST/RECV_NON_FIRST), the COMPLETE floor and the liveness stamp. + now_ns = time.monotonic_ns() + self._handle_response(resp, now_ns) if self._endpoint_response_idle_timeout_ns is not None: - self._last_response_progress_ns = time.monotonic_ns() + self._last_response_progress_ns = now_ns ``` Uses `recv()` exclusively — no `poll()` spin. The ZMQ fd is registered with the @@ -335,7 +338,7 @@ returns as soon as the fd is readable and `eager_task_factory` executes the woken semaphore waiter synchronously, there is no added latency compared to a poll-based approach. -**`_handle_response(resp)`**: +**`_handle_response(resp, now_ns)`**: - `QueryResult`: publish COMPLETE event, decrement `_inflight`, call `on_sample_complete`, call `strategy.on_query_complete(query_id)` if strategy @@ -352,12 +355,21 @@ poll-based approach. until the batch threshold is reached or `flush()` is called. The timestamp itself is still accurate (captured before buffering), but the EventRecord reaches subscribers with batching latency. +- RECV_FIRST and RECV_NON_FIRST: one `monotonic_ns()` read per response in + `_receive_responses`, when the session handles it. - COMPLETE: `QueryResult.completed_at` is set via `force_setattr(monotonic_ns())` in `__post_init__`, regenerated on - deserialization. Both ISSUED and COMPLETE timestamps share the same ZMQ - transit bias. TTFT (`RECV_FIRST - ISSUED`) is still sensitive to this overhead - since it spans the full ZMQ round-trip. TPOT avoids cross-process clock skew - by computing time deltas between consecutive chunks within the same process. + deserialization. The transport decodes a whole drain before the session + handles any message, so a result decoded in the same drain as its chunks is + older than their RECV stamps (negative TPOT). In that case only, COMPLETE + takes the handling stamp of the response's last chunk, so a response's + COMPLETE is never earlier than its own chunks. A response without chunks keeps + `completed_at`. When COMPLETE keeps `completed_at`, ISSUED and COMPLETE share + the same ZMQ transit bias. When it takes the last chunk's stamp, COMPLETE also + includes the session's handling backlog for that drain. TTFT + (`RECV_FIRST - ISSUED`) is still sensitive to this overhead since it spans the + full ZMQ round-trip. TPOT avoids cross-process clock skew by computing time + deltas between consecutive chunks within the same process. ### LoadStrategy (Protocol) @@ -647,7 +659,7 @@ sequenceDiagram end W-->>I: ZMQ PUSH (QueryResult — final accumulated output) I-->>B: recv() - B->>E: publish(COMPLETE, uuid, completed_at) + B->>E: publish(COMPLETE, uuid, max(completed_at, last chunk stamp)) end E->>M: ZMQ PUB (EventRecord) @@ -945,13 +957,20 @@ shorter strings, no parsing overhead. Update `Query.default_factory` to match. `issuer.issue(query)`. ZMQ push is sync/non-blocking — timestamp is honest. - **COMPLETE**: `QueryResult.completed_at` set in `__post_init__` on deserialization in the main process. Measures main-process receipt time, not - worker-side completion. + worker-side completion. If a chunk of the same response was handled later + (same transport drain), COMPLETE takes that chunk's stamp instead, so it is + never earlier than its own chunks. +- **RECV_FIRST / RECV_NON_FIRST**: one `monotonic_ns()` read per response in + `_receive_responses`, when the session handles it. Measures main-process + handling time. - **TTFT**: `RECV_FIRST - ISSUED` includes full round-trip ZMQ overhead (outbound to worker + return to main). This adds 20-100μs of systematic bias. Acceptable for most benchmarks; document as a known measurement overhead. - **Latency (COMPLETE - ISSUED)**: Both timestamps taken on the main process - side. ZMQ transit bias is symmetric and cancels. This is the most accurate - measurement. + side. When COMPLETE keeps `completed_at`, the ZMQ transit bias is symmetric + and cancels; this is the most accurate measurement. When COMPLETE is floored + to the response's last chunk, it also includes the session's handling backlog + for that drain. ### Stale Completions After Saturation @@ -959,8 +978,9 @@ After a warmup phase (no drain), in-flight responses arrive during the perf phase. The receiver must distinguish stale vs current-phase completions: ```python -def _handle_response(self, resp: QueryResult) -> None: +def _handle_response(self, resp: QueryResult, now_ns: int) -> None: query_id = resp.id + last_chunk_ns = self._last_chunk_ns.pop(query_id, None) # set by its chunks phase_issuer = self._current_phase_issuer # ERROR is published before COMPLETE (the aggregator depends on this order # for tracked-failure counting). @@ -971,7 +991,8 @@ def _handle_response(self, resp: QueryResult) -> None: if self._current_phase_type != PhaseType.WARMUP: self._publisher.publish(EventRecord( event_type=SampleEventType.COMPLETE, - timestamp_ns=resp.completed_at, + # completed_at, never earlier than this response's chunks + timestamp_ns=self._complete_timestamp_ns(resp, now_ns, last_chunk_ns), sample_uuid=query_id, )) # Only route to the current phase's strategy if this is a current-phase query. diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index 896646f81..bffc2492a 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/src/inference_endpoint/load_generator/session.py @@ -402,6 +402,10 @@ def __init__( self._drain_event = asyncio.Event() self._fatal_error: RuntimeError | None = None self._last_response_progress_ns: int | None = None + # Handling stamp of each streaming response's latest chunk, so that its + # COMPLETE is never earlier than its own chunks. Popped at the result; + # stamps of synthetically terminated queries are dropped at phase end. + self._last_chunk_ns: dict[str, int] = {} # Liveness guard: a single re-arming timer, not a task. It exists only # while work is in flight, so the disabled and idle paths cost nothing. self._endpoint_response_idle_timeout_ns = ( @@ -483,6 +487,7 @@ async def run( except asyncio.CancelledError: pass self._disarm_endpoint_response_idle_timeout() + self._last_chunk_ns.clear() if self._stop_requested: # Aborted run (Ctrl-C, transport closure, run watchdog): mark # it BEFORE the terminal ENDED so the aggregator's ENDED-driven @@ -578,6 +583,12 @@ async def _run_phase(self, phase: PhaseConfig) -> PhaseResult | None: if phase.drain_after: await self._drain_inflight(phase_issuer, phase.drain_timeout) + # A query terminated without a QueryResult (e.g. a turn timeout after + # some chunks) leaves a stamp that nothing else would pop. + for query_id in [ + q for q in self._last_chunk_ns if q in phase_issuer.completed_uuids + ]: + del self._last_chunk_ns[query_id] if phase.phase_type == PhaseType.PERFORMANCE: phase_issuer.stop_performance_tracking() @@ -750,6 +761,7 @@ def _handle_response(self, resp: QueryResult | StreamChunk, now_ns: int) -> None if isinstance(resp, QueryResult): query_id = resp.id + last_chunk_ns = self._last_chunk_ns.pop(query_id, None) # Drop late responses for queries already synthetically terminated # (e.g. by AgenticInferenceStrategy._handle_timeout). Without this gate, # a real response arriving after timeout double-publishes ERROR/COMPLETE @@ -793,9 +805,9 @@ def _handle_response(self, resp: QueryResult | StreamChunk, now_ns: int) -> None self._publisher.publish( EventRecord( event_type=SampleEventType.COMPLETE, - timestamp_ns=resp.completed_at - if isinstance(resp.completed_at, int) - else now_ns, + timestamp_ns=self._complete_timestamp_ns( + resp, now_ns, last_chunk_ns + ), sample_uuid=query_id, conversation_id=conv_id_str, turn=turn_num, @@ -833,6 +845,13 @@ def _handle_response(self, resp: QueryResult | StreamChunk, now_ns: int) -> None conv_id_str, turn_num = phase_issuer.uuid_to_conv_info.get( resp.id, ("", None) ) + if phase_issuer is None or ( + resp.id in phase_issuer.uuid_to_conv_info + and resp.id not in phase_issuer.completed_uuids + ): + self._last_chunk_ns[resp.id] = now_ns + else: + self._last_chunk_ns.pop(resp.id, None) self._publisher.publish( EventRecord( event_type=event_type, @@ -843,6 +862,26 @@ def _handle_response(self, resp: QueryResult | StreamChunk, now_ns: int) -> None ) ) + @staticmethod + def _complete_timestamp_ns( + resp: QueryResult, now_ns: int, last_chunk_ns: int | None + ) -> int: + """COMPLETE stamp: the decode time, but never before the response's chunks. + + `resp.completed_at` is set when the transport decodes the result. The + transport decodes a whole drain before the session handles any message, + so a result decoded in the same drain as its chunks is older than the + chunks' handling stamps (RECV_FIRST/RECV_NON_FIRST), which would give a + negative TPOT. Only that case moves COMPLETE, to the last chunk's stamp; + a response without chunks keeps its decode time. + """ + completed_ns = ( + resp.completed_at if isinstance(resp.completed_at, int) else now_ns + ) + if last_chunk_ns is None: + return completed_ns + return max(completed_ns, last_chunk_ns) + def _make_stop_check( self, settings: RuntimeSettings, phase_start_ns: int ) -> Callable[[], bool]: diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index 99cc04ced..87443a4ee 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -1200,6 +1200,57 @@ async def inject_streaming_responses(): recv_first = publisher.events_of_type(SampleEventType.RECV_FIRST) assert len(recv_first) == 2 + @pytest.mark.asyncio + async def test_streaming_complete_not_before_chunks_of_one_drain(self): + """COMPLETE is never stamped before the chunks of its own response. + + One transport drain decodes [chunk, chunk, result] before the session + handles any of them, so QueryResult.completed_at is older than the + chunks' handling time. Regression test for mlcommons/endpoints#533. + """ + loop = asyncio.get_running_loop() + publisher = FakePublisher() + + issuer = FakeIssuer() + issuer._loop = loop + issuer._auto_respond = False + + session = BenchmarkSession(issuer, publisher, loop) + settings = _make_settings( + load_pattern=LoadPattern( + type=LoadPatternType.CONCURRENCY, target_concurrency=1 + ), + n_samples=1, + ) + phases = [PhaseConfig("perf", settings, FakeDataset(1))] + + async def inject_one_drain(): + while not issuer._issued: + await asyncio.sleep(0.005) + query_id = issuer._issued[0].id + # Decoded first, as the transport decodes a whole drain at once. + result = QueryResult(id=query_id, response_output="out") + await asyncio.sleep(0.01) + issuer.inject_response( + StreamChunk(id=query_id, metadata={"first_chunk": True}) + ) + issuer.inject_response(StreamChunk(id=query_id, response_chunk="more")) + issuer.inject_response(result) + + injector = asyncio.create_task(inject_one_drain()) + await asyncio.wait_for(session.run(phases), timeout=5.0) + await injector + + recv_first = publisher.events_of_type(SampleEventType.RECV_FIRST) + recv_non_first = publisher.events_of_type(SampleEventType.RECV_NON_FIRST) + complete = publisher.events_of_type(SampleEventType.COMPLETE) + assert len(recv_first) == len(recv_non_first) == len(complete) == 1 + assert ( + recv_first[0].timestamp_ns + <= recv_non_first[0].timestamp_ns + <= complete[0].timestamp_ns + ) + @pytest.mark.asyncio async def test_concurrency_strategy_transport_close_no_deadlock(self): """ConcurrencyStrategy must not deadlock when transport closes mid-phase.""" @@ -1318,6 +1369,122 @@ async def inject_error(): f"complete at idx {complete_idx}" ) + @pytest.mark.asyncio + async def test_complete_stamp_is_decode_time_floored_by_own_chunks(self): + """COMPLETE keeps QueryResult.completed_at unless one of the response's + own chunks was handled later; then it takes that chunk's stamp.""" + loop = asyncio.get_running_loop() + publisher = FakePublisher() + session = BenchmarkSession(FakeIssuer(), publisher, loop) + + late = QueryResult(id="q-late", response_output="out") + session._handle_response( + StreamChunk(id="q-late", metadata={"first_chunk": True}), + late.completed_at + 500, + ) + session._handle_response(late, late.completed_at + 900) + early = QueryResult(id="q-early", response_output="out") + session._handle_response( + StreamChunk(id="q-early", metadata={"first_chunk": True}), + early.completed_at - 500, + ) + session._handle_response(early, early.completed_at + 900) + + complete = { + e.sample_uuid: e.timestamp_ns + for e in publisher.events_of_type(SampleEventType.COMPLETE) + } + assert complete == { + "q-late": late.completed_at + 500, + "q-early": early.completed_at, + } + assert session._last_chunk_ns == {} + + @pytest.mark.asyncio + async def test_chunk_stamp_kept_only_for_active_queries(self): + """Chunks of stale (previous-phase) or terminated queries may never get + a QueryResult, so their stamps are not kept.""" + loop = asyncio.get_running_loop() + issuer = FakeIssuer() + issuer._auto_respond = False + publisher = FakePublisher() + session = BenchmarkSession(issuer, publisher, loop) + phase_issuer = PhaseIssuer(FakeDataset(2), issuer, publisher, lambda: False) + session._current_phase_issuer = phase_issuer + active = phase_issuer.issue(0) + terminated = phase_issuer.issue(1) + phase_issuer.completed_uuids.add(terminated) + + for query_id in (active, terminated, "q-stale"): + session._handle_response( + StreamChunk(id=query_id, metadata={"first_chunk": True}), 100 + ) + + assert session._last_chunk_ns == {active: 100} + assert len(publisher.events_of_type(SampleEventType.RECV_FIRST)) == 3 + + @pytest.mark.asyncio + async def test_late_chunk_drops_stamp_of_terminated_query(self): + """A chunk that arrives after a synthetic termination removes the stamp + recorded before it, since no QueryResult will pop it.""" + loop = asyncio.get_running_loop() + issuer = FakeIssuer() + issuer._auto_respond = False + publisher = FakePublisher() + session = BenchmarkSession(issuer, publisher, loop) + phase_issuer = PhaseIssuer(FakeDataset(1), issuer, publisher, lambda: False) + session._current_phase_issuer = phase_issuer + query_id = phase_issuer.issue(0) + + session._handle_response( + StreamChunk(id=query_id, metadata={"first_chunk": True}), 100 + ) + phase_issuer.completed_uuids.add(query_id) + phase_issuer.uuid_to_conv_info.pop(query_id) + session._handle_response(StreamChunk(id=query_id, response_chunk="x"), 200) + + assert session._last_chunk_ns == {} + + @pytest.mark.asyncio + async def test_phase_end_drops_chunk_stamps_of_terminated_queries(self): + """A turn timeout after some chunks publishes a synthetic failure and no + QueryResult follows, so the phase end drops that query's stamp.""" + loop = asyncio.get_running_loop() + issuer = FakeIssuer() + issuer._auto_respond = False + issuer._loop = loop + session = BenchmarkSession(issuer, FakePublisher(), loop) + + class TimeoutAfterChunkStrategy: + async def execute(self, phase_issuer) -> None: + query_id = phase_issuer.issue(0) + issuer.inject_response( + StreamChunk(id=query_id, metadata={"first_chunk": True}) + ) + while query_id not in session._last_chunk_ns: + await asyncio.sleep(0) + # The state changes of AgenticInferenceStrategy._handle_timeout. + phase_issuer.mark_inflight_complete() + phase_issuer.completed_uuids.add(query_id) + phase_issuer.uuid_to_conv_info.pop(query_id) + + def on_query_complete(self, query_id: str) -> None: + pass + + phase = PhaseConfig( + "perf", + _make_settings(n_samples=1), + FakeDataset(1), + strategy=TimeoutAfterChunkStrategy(), + ) + session._recv_task = asyncio.create_task(session._receive_responses()) + try: + await asyncio.wait_for(session._run_phase(phase), timeout=1.0) + assert session._last_chunk_ns == {} + finally: + issuer.shutdown() + await session._recv_task + @pytest.mark.asyncio async def test_handle_response_stamps_conversation_id_and_turn(self): """All event types inherit (conv_id, turn) seeded at issue time; streaming @@ -1371,10 +1538,8 @@ async def test_handle_response_stamps_conversation_id_and_turn(self): assert [(e.conversation_id, e.turn) for e in complete] == [("conv-9", 5)] assert complete[0].finish_reason == "stop" assert complete[0].worker_id == 3 - # QueryResult stamps completed_at itself (force_setattr in __post_init__), - # and that stamp outranks the caller's arrival timestamp. + # Without chunks, COMPLETE keeps the decode-time QueryResult.completed_at. assert complete[0].timestamp_ns == ok_resp.completed_at - assert complete[0].timestamp_ns != 1_002 assert "q-ok" not in phase_issuer.uuid_to_conv_info assert "q-ok" not in phase_issuer.completed_uuids