From 25563bbfc986ba5636fe050c11010c84921a58ee Mon Sep 17 00:00:00 2001 From: Calin Teodor <45018530+TeoSlayer@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:04:02 -0700 Subject: [PATCH 1/2] fix: retain session start in metrics reports --- .../services/metrics_aggregator/aggregator.py | 1 + .../services/metrics_aggregator/publisher.py | 9 ++++++- .../services/metrics_aggregator/registry.py | 2 ++ .../services/metrics_aggregator/snapshot.py | 3 +++ src/inference_endpoint/metrics/report.py | 2 +- .../metrics_aggregator/test_aggregator.py | 3 ++- .../metrics_aggregator/test_publisher.py | 25 ++++++++++++++++--- .../metrics_aggregator/test_snapshot.py | 8 +++++- tests/unit/metrics/test_report_builder.py | 13 ++++++++++ 9 files changed, 59 insertions(+), 7 deletions(-) diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py index 6bab4609e..2a202648c 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py @@ -426,6 +426,7 @@ async def process(self, records: list[EventRecord]) -> None: self._session_state, self.pending_tokens, ), + session_started_ns=self._session_start_ns, ) if ( self._token_queue is not None diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/publisher.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/publisher.py index 26167f771..0103c1593 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/publisher.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/publisher.py @@ -87,6 +87,7 @@ def __init__( self._loop = loop self._final_snapshot_path = final_snapshot_path self._tick_task: asyncio.Task | None = None + self._session_started_ns: int | None = None self._closed = False # publish_final is idempotent AND serialized: the SIGTERM handler in # __main__.py and the aggregator's ENDED-driven path can both call @@ -107,6 +108,8 @@ def start( registry: MetricsRegistry, publish_interval_s: float, get_runtime_state: Callable[[], tuple[SessionState, int]], + *, + session_started_ns: int | None = None, ) -> None: """Begin publishing live ticks every ``publish_interval_s`` seconds. @@ -133,6 +136,7 @@ def start( raise ValueError( f"publish_interval_s must be positive, got {publish_interval_s}" ) + self._session_started_ns = session_started_ns async def _tick() -> None: while True: @@ -140,7 +144,9 @@ async def _tick() -> None: await asyncio.sleep(publish_interval_s) state, n_pending = get_runtime_state() snap = registry.build_snapshot( - state=state, n_pending_tasks=n_pending + state=state, + n_pending_tasks=n_pending, + session_started_ns=self._session_started_ns, ) self._publisher.publish(snap) except asyncio.CancelledError: @@ -237,6 +243,7 @@ async def _publish_final_locked( state=terminal_state, n_pending_tasks=n_pending_tasks, steady_state=steady_state, + session_started_ns=self._session_started_ns, ) # Primary: atomic JSON file write. Run on a worker thread because diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py index 18f874ae9..927cd99f8 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py @@ -512,6 +512,7 @@ def build_snapshot( state: SessionState, n_pending_tasks: int, steady_state: dict | None = None, + session_started_ns: int | None = None, ) -> MetricsSnapshot: # Exact (raw-values) computation is reserved for the COMPLETE snapshot; # live and draining snapshots use the cheap HDR path. @@ -529,6 +530,7 @@ def build_snapshot( n_pending_tasks=n_pending_tasks, metrics=metrics, steady_state=steady_state, + session_started_ns=session_started_ns, ) # -- introspection (mostly for tests) -------------------------------- diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/snapshot.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/snapshot.py index f32c59033..6f417bc08 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/snapshot.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/snapshot.py @@ -171,6 +171,7 @@ class MetricsSnapshot( ``metrics/steady_state_diagnostics.py``. Present only on the terminal snapshot when collection was enabled and the series describe the run. + session_started_ns: Monotonic session start; None before STARTED. ``array_like=True`` makes field order the wire format. Append fields only. """ @@ -181,6 +182,7 @@ class MetricsSnapshot( n_pending_tasks: int metrics: list[MetricStat] steady_state: SteadyState | None = None + session_started_ns: int | None = None # 4-byte topic to match TOPIC_FRAME_SIZE-prefix protocol used by the @@ -249,6 +251,7 @@ def snapshot_to_dict(snap: MetricsSnapshot) -> dict: return { "counter": snap.counter, "timestamp_ns": snap.timestamp_ns, + "session_started_ns": snap.session_started_ns, "state": snap.state.value, "n_pending_tasks": snap.n_pending_tasks, "metrics": [_metric_to_dict(m) for m in snap.metrics], diff --git a/src/inference_endpoint/metrics/report.py b/src/inference_endpoint/metrics/report.py index 98834de34..661a16729 100644 --- a/src/inference_endpoint/metrics/report.py +++ b/src/inference_endpoint/metrics/report.py @@ -429,7 +429,7 @@ def _series_dict(key: str) -> dict[str, Any]: return cls( version=str(version_info.get("version", "unknown")), git_sha=version_info.get("git_sha"), - test_started_at=0, # TODO: surface session_started_ns via snapshot + test_started_at=snap.get("session_started_ns") or 0, n_samples_issued=n_issued, n_samples_completed=n_completed, n_samples_failed=n_failed, diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py index b2d5f1b64..f2e55598b 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py @@ -741,10 +741,11 @@ async def test_error_event_increments_total_failed(self, tmp_path): async def test_session_started_stores_timestamp(self, tmp_path): loop = asyncio.get_event_loop() with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: - agg, _, _ = make_aggregator(ctx, loop, "agg_started_ts") + agg, _, publisher = make_aggregator(ctx, loop, "agg_started_ts") try: await agg.process([session_event(SessionEventType.STARTED, ts=42)]) assert agg._table.session_started_ns == 42 + assert publisher.start.call_args.kwargs["session_started_ns"] == 42 finally: agg.close() diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_publisher.py b/tests/unit/async_utils/services/metrics_aggregator/test_publisher.py index 9e26f734a..8e61a01c2 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_publisher.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_publisher.py @@ -67,10 +67,13 @@ def get_runtime_state() -> tuple[SessionState, int]: calls.append(True) return SessionState.LIVE, 0 + inner_mock = MagicMock() + publisher._publisher = inner_mock publisher.start( registry, publish_interval_s=0.01, get_runtime_state=get_runtime_state, + session_started_ns=42, ) assert publisher._tick_task is not None assert not publisher._tick_task.done() @@ -78,12 +81,15 @@ def get_runtime_state() -> tuple[SessionState, int]: # Let at least one tick run. await asyncio.sleep(0.05) assert len(calls) >= 1 + assert inner_mock.publish.call_args.args[0].session_started_ns == 42 finally: publisher.close() + @pytest.mark.parametrize("started", [False, True]) + @pytest.mark.parametrize("interrupted", [False, True]) @pytest.mark.asyncio async def test_publish_final_writes_json_atomically( - self, tmp_path: Path, zmq_ctx_scope: ManagedZMQContext + self, tmp_path: Path, zmq_ctx_scope: ManagedZMQContext, started, interrupted ): loop = asyncio.get_event_loop() target = tmp_path / "final_snapshot.json" @@ -99,7 +105,16 @@ async def test_publish_final_writes_json_atomically( registry.register_counter("c") registry.increment("c", 5) - await publisher.publish_final(registry, n_pending_tasks=0) + if started: + publisher.start( + registry, + publish_interval_s=1, + get_runtime_state=lambda: (SessionState.LIVE, 0), + session_started_ns=42, + ) + await publisher.publish_final( + registry, n_pending_tasks=0, interrupted=interrupted + ) # The .tmp file MUST NOT exist after the rename. tmp_target = target.with_suffix(target.suffix + ".tmp") @@ -107,8 +122,12 @@ async def test_publish_final_writes_json_atomically( assert target.exists(), "final snapshot should be on disk" decoded = json.loads(target.read_bytes()) - assert decoded["state"] == SessionState.COMPLETE.value + expected_state = ( + SessionState.INTERRUPTED if interrupted else SessionState.COMPLETE + ) + assert decoded["state"] == expected_state.value assert decoded["n_pending_tasks"] == 0 + assert decoded["session_started_ns"] == (42 if started else None) finally: publisher.close() diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_snapshot.py b/tests/unit/async_utils/services/metrics_aggregator/test_snapshot.py index 7fb304278..81249a531 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_snapshot.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_snapshot.py @@ -80,13 +80,15 @@ def test_roundtrip(self, sum_sq): @pytest.mark.unit class TestMetricsSnapshot: - def test_empty_metrics_roundtrip(self): + @pytest.mark.parametrize("started_ns", [None, 0, 42]) + def test_empty_metrics_roundtrip(self, started_ns): snap = MetricsSnapshot( counter=1, timestamp_ns=1234, state=SessionState.LIVE, n_pending_tasks=0, metrics=[], + session_started_ns=started_ns, ) codec = MetricsSnapshotCodec() topic, payload = codec.encode(snap) @@ -95,6 +97,10 @@ def test_empty_metrics_roundtrip(self): decoded = codec.decode(payload) assert decoded == snap + def test_decode_without_session_start(self): + payload = msgspec.msgpack.encode([1, 1234, "live", 0, [], None]) + assert MetricsSnapshotCodec().decode(payload).session_started_ns is None + def test_tagged_union_dispatch(self): """Decoder must produce the right concrete type per tag.""" snap = MetricsSnapshot( diff --git a/tests/unit/metrics/test_report_builder.py b/tests/unit/metrics/test_report_builder.py index b284f983d..271520547 100644 --- a/tests/unit/metrics/test_report_builder.py +++ b/tests/unit/metrics/test_report_builder.py @@ -383,6 +383,19 @@ def test_complete_flag_false_when_state_live(self): @pytest.mark.unit class TestReportDisplayAndSerialize: + @pytest.mark.parametrize( + "start", [{}, {"session_started_ns": None}, {"session_started_ns": 42}] + ) + def test_start_time_from_snapshot(self, start): + snap = snapshot_to_dict( + _make_registry().build_snapshot( + state=SessionState.COMPLETE, n_pending_tasks=0 + ) + ) + snap.update(start) + report = Report.from_snapshot(snap) + assert report.test_started_at == (start.get("session_started_ns") or 0) + def test_display_summary(self): registry = _make_registry(n_samples=10) report = _build_report(registry) From dc3bd25f42c84f6fc38cb1df0062a888ff81939b Mon Sep 17 00:00:00 2001 From: Calin Teodor <45018530+TeoSlayer@users.noreply.github.com> Date: Thu, 5 Nov 2026 09:38:38 -0800 Subject: [PATCH 2/2] test: wait for benchmark startup before wedging metrics --- tests/integration/commands/test_sigint.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/integration/commands/test_sigint.py b/tests/integration/commands/test_sigint.py index 1197744dc..d6986981e 100644 --- a/tests/integration/commands/test_sigint.py +++ b/tests/integration/commands/test_sigint.py @@ -276,6 +276,7 @@ def test_sigint_grace_expiry_abandons_wedged_drain(mock_http_echo_server, tmp_pa """ report_dir = tmp_path / "report" config_path = tmp_path / "bench.yaml" + log_file = tmp_path / "run.log" _write_config( report_dir, mock_http_echo_server.url, @@ -286,10 +287,18 @@ def test_sigint_grace_expiry_abandons_wedged_drain(mock_http_echo_server, tmp_pa agg_pid: int | None = None try: with _benchmark_proc( - [_cli(), "benchmark", "from-config", "-c", str(config_path)] + [_cli(), "benchmark", "from-config", "-c", str(config_path)], + log_file, ) as proc: _wait_services_ready(proc, report_dir) - time.sleep(3.0) # comfortably inside the ~120 s performance phase + # Service readiness precedes client startup. Wait for the session + # before allowing live metrics ticks to reach the parent. + deadline = time.monotonic() + 60.0 + while "Starting phase:" not in log_file.read_text(errors="replace"): + assert proc.poll() is None, log_file.read_text(errors="replace") + assert time.monotonic() < deadline, "performance phase never started" + time.sleep(0.05) + time.sleep(3.0) agg_pid = _pid_of_child("metrics_aggregator", str(report_dir)) assert agg_pid is not None, "aggregator child not found"