Change MetricsAggregatorService to publisher to support live metrics without shmem

[nv-alicheng]

What does this PR do?

shmem implementation of KVStore in MetricsAggregatorService causes issues on ARM. Several solutions exist:

1. Rewrite shmem implementation in C/C++ where memory fencing primitives are exposed.
2. Restructure the design of the MetricsAggregatorService to be a ZMQ Publisher which publishes metrics at a fixed rate, which other processes, such as a TUI, can subscribe to.

This PR implements (2).

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor/cleanup

Related issues

Testing

• Tests added/updated
• All tests pass locally
• Manual testing completed

Checklist

• Code follows project style
• Pre-commit hooks pass
• Documentation updated (if needed)

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[gemini-code-assist]

Code Review

This pull request refactors the metrics aggregation system to use a registry-based architecture with HDR histograms and a generalized pub/sub transport layer, replacing the legacy mmap-backed storage. The update introduces periodic snapshot publishing with disk fallback and updates reporting logic to consume these snapshots. Feedback suggests improving encapsulation by exposing in-flight task metrics through public properties and adopting a more numerically stable variance formula for high-precision latency calculations. Additionally, several legacy tests have been skipped pending migration to the new system.

[gemini-code-assist]

Directly accessing the protected member _in_flight_tasks of MetricsTable breaks encapsulation. It is better to expose this information through a public property on the MetricsTable class.

Suggested change

	len(table._in_flight_tasks),
	len(table.in_flight_tasks_count),

[gemini-code-assist]

Directly accessing the protected member _in_flight_tasks of MetricsTable breaks encapsulation. It is better to expose this information through a public property on the MetricsTable class.

Suggested change

	logger.info("Draining %d async tasks...", len(table._in_flight_tasks))
	logger.info("Draining %d async tasks...", table.in_flight_tasks_count)

[gemini-code-assist]

Iterating over and cancelling tasks directly from the aggregator violates the encapsulation of MetricsTable. This logic should be moved into a dedicated method within MetricsTable.

                table.cancel_all_tasks()

[gemini-code-assist]

Directly accessing the protected member _in_flight_tasks of MetricsTable breaks encapsulation. It is better to expose this information through a public property on the MetricsTable class.

Suggested change

	n_pending = len(table._in_flight_tasks)
	n_pending = table.in_flight_tasks_count

[gemini-code-assist]

To improve encapsulation and support the refactored MetricsAggregatorService, please expose the count of in-flight tasks and a method to cancel them.

        self._in_flight_tasks: set[asyncio.Task] = set()

    @property
    def in_flight_tasks_count(self) -> int:
        """Number of async trigger tasks currently in flight."""
        return len(self._in_flight_tasks)

    def cancel_all_tasks(self) -> None:
        """Cancel all in-flight async trigger tasks."""
        for task in self._in_flight_tasks:
            if not task.done():
                task.cancel()

[gemini-code-assist]

The naive variance formula sum_sq - total^2 / n is numerically unstable when total^2 / n is very close to sum_sq, which is common for high-precision latency metrics in nanoseconds. Since total and sum_sq are Python integers for latency series, using the integer-based formula (n * sum_sq - total^2) / (n * (n - 1)) provides exact precision before the final square root.

Suggested change

	var_num = stat.sum_sq - stat.total * stat.total / n
	std_dev = math.sqrt(var_num / (n - 1)) if var_num > 0 else 0.0
	# Use a more numerically stable formula for variance when aggregates are large.
	# If using integers, (n * sum_sq - total^2) is exact.
	if isinstance(stat.total, int) and isinstance(stat.sum_sq, int):
	var_num = (n * stat.sum_sq - stat.total**2)
	std_dev = math.sqrt(max(0, var_num)) / math.sqrt(n * (n - 1))
	else:
	var_num = stat.sum_sq - stat.total * stat.total / n
	std_dev = math.sqrt(max(0, var_num / (n - 1)))

[arekay-nv]

(superseded — see #pullrequestreview-4237753512 for the full review-council output)

[arekay-nv]

Review Council — Multi-AI Code Review

Reviewed by: Codex + Claude | Depth: thorough

Found 12 issues across 7 files (excluding tests, pyproject.toml, uv.lock).

⚠️ GitHub's review-comment API is currently returning a persistent pull_request_review_thread.base internal error on this PR for inline comments via line+side and position-based posts (the github-code-quality and gemini-code-assist bot comments from yesterday were unaffected). Falling back to a single review with file:line links — clicking each link opens the PR diff at the cited line.

Note: the existing gemini comments on aggregator.py:286/342/353/354 already cover the encapsulation concern around table._in_flight_tasks direct access, so that issue is intentionally not duplicated here. The cancel-without-await issue (#4 below) is a separate concurrency concern at the same file:line.

🔴 Must Fix (critical/high)

Issues that will cause incorrect behavior, data loss, or security problems in production.

#	File:Line	Reviewer	Category	Summary
1	src/inference_endpoint/commands/benchmark/execute.py:482	Claude	data-integrity	Subscriber late-binding can drop initial ticks (ZMQ slow-joiner)
2	src/inference_endpoint/async_utils/services/metrics_aggregator/publisher.py:168	Claude	performance	_write_atomic_fallback runs blocking I/O on the event loop
3	src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py:208	Codex	performance	Unbounded raw-sample retention in SeriesSampler

🟡 Should Fix (medium)

Real issues that trigger under specific conditions or design flaws that will compound.

#	File:Line	Reviewer	Category	Summary
4	src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py:353	Claude	error-handling	Cancellations not awaited before reading n_pending
5	src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py:161	Claude	bug	HdrHistogram(low, high, sig_figs) constructed without high >= 2*low check
6	src/inference_endpoint/async_utils/services/metrics_aggregator/subscriber.py:55	Claude	data-integrity	conflate=True default for the Report consumer is fragile
7	src/inference_endpoint/commands/benchmark/execute.py:423	Codex	data-integrity	Stale final_snapshot.msgpack not cleared on report-dir reuse

🔵 Consider (low)

Valid improvements that could be follow-ups.

#	File:Line	Reviewer	Category	Summary
8	src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py:281	Claude	bug	Double STARTED reassigns and orphans the tick task
9	src/inference_endpoint/async_utils/services/metrics_aggregator/publisher.py:195	Claude	concurrency	close() cancels tick task but doesn't await it
10	src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py:166	Claude	error-handling	No top-level exception handling around run_until_complete(main())
11	src/inference_endpoint/load_generator/session.py:408	Claude	design	ERROR-before-COMPLETE relies on an undocumented publisher ordering contract
12	src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py:132	Claude	bug	SIGTERM bypasses publish_final

---

Detailed findings

1. 🔴 high · data-integrity · [Claude]

src/inference_endpoint/commands/benchmark/execute.py:482

Subscriber late-binding can drop initial ticks (ZMQ slow-joiner)

MetricsSnapshotSubscriber is constructed (482) and start()ed (485) AFTER launcher.launch() waits for aggregator readiness. The aggregator can begin publishing as soon as STARTED is observed; ZMQ slow-joiner means snapshots emitted before the SUB handshake completes are dropped. The comment on line 480 acknowledges this. On a busy host the subscriber may also miss COMPLETE if the publisher closes before the subscription warms up, silently degrading to LIVE/DRAINING (incomplete report). Fix: move subscriber construction + start() BEFORE launcher.launch(). Connecting to a not-yet-bound IPC path is fine — ZMQ retries connect transparently.

2. 🔴 high · performance · [Claude]

src/inference_endpoint/async_utils/services/metrics_aggregator/publisher.py:168

_write_atomic_fallback runs blocking I/O on the event loop

f.write + f.flush + two os.fsync (file + parent dir) + os.rename execute synchronously on the aggregator subprocess's main event loop (awaited via publish_final from aggregator.py:362). On a busy box os.fsync can block tens-to-hundreds of ms — long enough to back-pressure event-record processing if any events are still in flight at finalization. Fix: wrap with await asyncio.to_thread(self._write_atomic_fallback, payload).

3. 🔴 high · performance · [Codex]

src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py:208

Unbounded raw-sample retention in SeriesSampler

self._raw.append(value) retains every observed value in an array.array for the full run so the final snapshot can recompute exact percentiles. The prior mmap-backed store spilled this to disk; the new path keeps it in the aggregator subprocess RAM, so memory now scales linearly with run length. A 5-min 50k-QPS run produces ~15M samples (~120MB per int64 series); once latency/ISL/OSL/TTFT/TPOT are all tracked, several hundred MB; longer runs OOM or swap heavily before the report is emitted. Fix: periodically spill _raw to disk, cap with reservoir sampling, or rely solely on the HDR digest for percentiles.

4. 🟡 medium · error-handling · [Claude]

src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py:353

Cancellations not awaited before reading n_pending

After the drain timeout, t.cancel() only schedules cancellation at the next await point. n_pending = len(table._in_flight_tasks) on the next line therefore reads a count that's still high, and the not-yet-cancelled tasks may still be running when _publisher.close() and the loop teardown happen → "Task was destroyed but it is pending!" warnings on shutdown. Fix: after t.cancel(), do await asyncio.gather(*tasks, return_exceptions=True) (bounded by a short timeout) before recomputing n_pending. Independent of the encapsulation point gemini already raised.

5. 🟡 medium · bug · [Claude]

src/inference_endpoint/async_utils/services/metrics_aggregator/registry.py:161

HdrHistogram(low, high, sig_figs) constructed without high >= 2*low check

hdrh requires high >= 2*low. Current bounds (1ns–3.6e12ns, 1–1e7 tokens) satisfy this, but a future caller passing e.g. hdr_low=1, hdr_high=1 (or any hdr_high < 2*hdr_low) gets an opaque ValueError from deep inside the C library. Fix: explicit pre-check after the clamps: if self._hdr_high < self._hdr_low * 2: raise ValueError(...).

6. 🟡 medium · data-integrity · [Claude]

src/inference_endpoint/async_utils/services/metrics_aggregator/subscriber.py:55

conflate=True default for the Report consumer is fragile

The docstring argues this is safe because COMPLETE is the publisher's last message, but ZMQ CONFLATE only keeps the most recent unread message. If the subscriber's loop is starved (main process busy in wait_for_exit) and the publisher closes before the SUB sees COMPLETE (LINGER expires, IPC socket unlinked), complete never fires and latest may be a stale LIVE. The 2 s wait_for_complete timeout in execute.py:548 papers over most cases, but consider conflate=False for the Report consumer (a single producer at a few snapshots/sec, not a TUI).

7. 🟡 medium · data-integrity · [Codex]

src/inference_endpoint/commands/benchmark/execute.py:423

Stale final_snapshot.msgpack not cleared on report-dir reuse

metrics_output_dir = ctx.report_dir / "metrics" is created (423) but an existing final_snapshot.msgpack from a prior run is never removed. If a user reuses --report-dir and the new run exits before publishing a fresh final snapshot (SIGTERM, crash), _load_final_snapshot_from_disk() decodes the previous run's file and silently builds the new report from stale metrics. Fix: unlink any pre-existing final_snapshot.msgpack here, OR refuse to decode a snapshot whose timestamp/run-id doesn't match the current run.

8. 🔵 low · bug · [Claude]

src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py:281

Double STARTED reassigns and orphans the tick task

self._publisher.start(...) is called every time a SessionEventType.STARTED record is observed. If two STARTED events ever land in the EventRecord stream (replay, buggy producer, test fixture), start() reassigns _tick_task; the previous task keeps running until garbage-collected and races with the new task to publish snapshots. Fix: guard with if self._tick_task is None: ... inside publisher.start(), or only call start() on the first STARTED in process().

9. 🔵 low · concurrency · [Claude]

src/inference_endpoint/async_utils/services/metrics_aggregator/publisher.py:195

close() cancels tick task but doesn't await it

self._tick_task.cancel() (205) is followed immediately by self._publisher.close() (206). If close() is invoked from a sync path, the cancelled task may still be runnable when the underlying ZMQ publisher socket is closed → brief window where _tick could call publish() on a closed transport. ZmqMessagePublisher.publish() early-returns on is_closed so this isn't fatal, but produces noisy CancelledError-during-shutdown traces. Fix: make close() async-aware (mirror publish_final) or document that callers must await publish_final (which already cancels the task) before close.

10. 🔵 low · error-handling · [Claude]

src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py:166

No top-level exception handling around run_until_complete(main())

If main() raises (bad CLI args, ZMQ bind failure, tokenizer load error), the subprocess exits with a stack trace but the parent's ServiceLauncher.wait_for_exit sees only a non-zero exit code with no diagnostic context propagated upward. Fix: wrap with try/except + structured logging so failures are surfaced in the parent's logs alongside the exit code.

11. 🔵 low · design · [Claude]

src/inference_endpoint/load_generator/session.py:408

ERROR-before-COMPLETE relies on an undocumented publisher ordering contract

The aggregator's TRACKED_SAMPLES_FAILED logic (aggregator.py:303–305) assumes ERROR is observed BEFORE COMPLETE removes the row. The ordering is correct here, but it depends on ZmqMessagePublisher preserving publish() order through batching AND ZMQ PUB→SUB delivering in order to a single SUB (both currently true). A short comment in session.py noting "publisher must preserve publish-order" would harden this against future transport refactors.

12. 🔵 low · bug · [Claude]

src/inference_endpoint/async_utils/services/metrics_aggregator/aggregator.py:132

SIGTERM bypasses publish_final

_finalize() only sets _shutdown_event after publish_final and _publisher.close() complete, but __main__.py doesn't install a SIGTERM handler. On parent-side launcher.kill_all (or any external SIGTERM before ENDED arrives), neither pub/sub COMPLETE nor disk fallback runs, leaving the parent's triple-redundant snapshot path empty. Fix: loop.add_signal_handler(SIGTERM, ...) to flush the final snapshot defensively before exit.

---

Generated by /review-council — Codex gpt-5.4 (review of git diff against main) + Claude (direct file review with HEAD-source line verification).