fix(session): stamp COMPLETE with the arrival time of its chunks

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,22 +957,30 @@ 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
 
 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.

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]:

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