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 event records
+        # (RECV_FIRST/RECV_NON_FIRST/COMPLETE) 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,16 @@ 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.
-- 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.
+- RECV_FIRST, RECV_NON_FIRST and COMPLETE: one `monotonic_ns()` read per
+  response in `_receive_responses`, when the session handles it. All three use
+  this same clock point, so a response's COMPLETE is never earlier than its own
+  chunks. (`QueryResult.completed_at` is set at decode; the transport decodes a
+  whole drain before the session handles any message, so a result decoded in the
+  same drain as its chunks would precede them and give a negative TPOT.) 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.
 
 ### LoadStrategy (Protocol)
 
@@ -635,7 +642,7 @@ sequenceDiagram
     alt Non-streaming
         W-->>I: ZMQ PUSH (QueryResult)
         I-->>B: recv()
-        B->>E: publish(COMPLETE, uuid, completed_at)
+        B->>E: publish(COMPLETE, uuid, timestamp_ns)
     else Streaming
         W-->>I: ZMQ PUSH (StreamChunk first_chunk)
         I-->>B: recv()
@@ -647,7 +654,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, timestamp_ns)
     end
 
     E->>M: ZMQ PUB (EventRecord)
@@ -943,9 +950,11 @@ shorter strings, no parsing overhead. Update `Query.default_factory` to match.
 
 - **ISSUED**: `monotonic_ns()` taken in `PhaseIssuer.issue()` immediately before
   `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.
+- **RECV_FIRST / RECV_NON_FIRST / COMPLETE**: one `monotonic_ns()` read per
+  response in `_receive_responses`, when the session handles it. Measures
+  main-process handling time, not worker-side completion. All three use the same
+  clock point, so a response's COMPLETE is never earlier than its own chunks.
+  `QueryResult.completed_at` is set at decode and does not give this guarantee.
 - **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.
@@ -959,7 +968,7 @@ 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
     phase_issuer = self._current_phase_issuer
     # ERROR is published before COMPLETE (the aggregator depends on this order
@@ -971,7 +980,7 @@ 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,
+            timestamp_ns=now_ns,  # the caller's per-response arrival stamp
             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

@@ -793,9 +793,11 @@ 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,
+                        # The same arrival stamp as RECV_FIRST/RECV_NON_FIRST.
+                        # resp.completed_at is set at decode, so a result decoded
+                        # in the same transport drain as its chunks would precede
+                        # them (negative TPOT).
+                        timestamp_ns=now_ns,
                         sample_uuid=query_id,
                         conversation_id=conv_id_str,
                         turn=turn_num,

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."""
@@ -1371,10 +1422,9 @@ 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.
-        assert complete[0].timestamp_ns == ok_resp.completed_at
-        assert complete[0].timestamp_ns != 1_002
+        # COMPLETE carries the caller's arrival stamp, like the chunk events,
+        # not the decode-time QueryResult.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