fix(metrics): drain pending count, SIGTERM task GC, NaN display

Addresses the three high-priority findings from the review council:

H1: drain_tasks now owns the timeout + cancel-and-await sequence, so the
pending count is captured before per-task done callbacks empty the
in-flight set. Previously read 0 unconditionally — the documented
state==COMPLETE and n_pending_tasks>0 drain-timeout contract was
unenforceable.

H2: Extract _make_sigterm_handler returning a strong-ref set[Task] that
holds the spawned _signal_finalize task; the loop tracks tasks via
weakref only, so a discarded create_task() return value can be GC'd
mid-flight (Python asyncio docs) — exactly the failure the INTERRUPTED
delivery path exists to prevent.

H3: _scrub_nonfinite maps producer-side NaN/Inf to None for strict JSON.
_display_metric did val * scale_factor with no guard → TypeError on
display(), which finalize_benchmark calls outside the report-build
try/except. Render N/A for None across named scalars, histogram
bucket edges, and percentiles.

Tests added (all verified failing pre-fix):
- test_drain_timeout_reports_pending_count: forever-blocking pool +
  drain_timeout_s=0.05, asserts publish_final receives n_pending>0
- test_sigterm_handler_holds_strong_reference_to_finalize_task: drives
  the handler, asserts task is in the strong-ref set, survives
  gc.collect(), and self-removes via done-callback on completion
- test_sigterm_handler_refreshes_tracked_duration: handler mirrors the
  ENDED path's tracked_duration_ns refresh before publish_final
- test_display_handles_scrubbed_nan_percentiles: dict with scrubbed
  None percentile values does not crash display(); renders N/A
- test_scrub_nonfinite_round_trip_yields_none: registry-side NaN/Inf
  surfaces as None in snapshot_to_dict and round-trips through
  json.dumps(allow_nan=False)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: nv-alicheng

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

@@ -19,6 +19,7 @@
 import asyncio
 import logging
 import signal
+from collections.abc import Callable
 from contextlib import AbstractContextManager, nullcontext
 from pathlib import Path
 
@@ -28,6 +29,7 @@
 from inference_endpoint.utils.logging import setup_logging
 
 from .aggregator import MetricCounterKey, MetricsAggregatorService
+from .metrics_table import MetricsTable
 from .publisher import MetricsPublisher
 from .registry import MetricsRegistry
 from .snapshot import MetricsSnapshotCodec
@@ -36,6 +38,63 @@
 logger = logging.getLogger(__name__)
 
 
+def _make_sigterm_handler(
+    *,
+    loop: asyncio.AbstractEventLoop,
+    registry: MetricsRegistry,
+    publisher: MetricsPublisher,
+    table: MetricsTable,
+    shutdown_event: asyncio.Event,
+) -> tuple[Callable[[], None], set[asyncio.Task]]:
+    """Build the SIGTERM handler that writes the INTERRUPTED final snapshot.
+
+    Returns ``(handler, pending_tasks)``. ``pending_tasks`` is the
+    strong-reference container that keeps spawned finalize tasks alive
+    while they run: asyncio tracks tasks only by weakref, so a task
+    whose only reference is the local variable inside the handler can
+    be garbage-collected mid-execution (per Python's asyncio docs).
+    Each spawned task self-removes from the set via
+    ``add_done_callback`` once it completes.
+
+    Exposed at module level (rather than nested in ``main()``) so the
+    GC-safety contract is unit-testable without driving the whole
+    subprocess lifecycle.
+    """
+    pending_tasks: set[asyncio.Task] = set()
+
+    async def _signal_finalize() -> None:
+        try:
+            # Mirror the ENDED-driven path: refresh tracked_duration_ns
+            # from the table BEFORE publish_final, otherwise an
+            # interrupted run whose STOP_PERFORMANCE_TRACKING never
+            # fired would report duration_ns=0 and QPS=N/A in the final
+            # report even after processing many tracked samples.
+            registry.set_counter(
+                MetricCounterKey.TRACKED_DURATION_NS.value,
+                table.total_tracked_duration_ns,
+            )
+            await publisher.publish_final(
+                registry,
+                n_pending_tasks=table.in_flight_tasks_count,
+                interrupted=True,
+            )
+        except Exception:  # noqa: BLE001 — best-effort.
+            logger.exception(
+                "metrics aggregator: SIGTERM-triggered publish_final failed"
+            )
+        shutdown_event.set()
+
+    def _on_sigterm() -> None:
+        logger.warning(
+            "metrics aggregator received SIGTERM; " "writing INTERRUPTED final snapshot"
+        )
+        task = loop.create_task(_signal_finalize())
+        pending_tasks.add(task)
+        task.add_done_callback(pending_tasks.discard)
+
+    return _on_sigterm, pending_tasks
+
+
 async def main() -> None:
     parser = argparse.ArgumentParser(
         description="Metrics aggregator service - subscribes to EventRecords and computes real-time metrics"
@@ -206,37 +265,14 @@ async def main() -> None:
             # aggregator's finalize — so we install a no-op handler for
             # SIGINT here, which prevents Python's default
             # KeyboardInterrupt and lets the parent control the lifecycle.
-            def _on_sigterm() -> None:
-                logger.warning(
-                    "metrics aggregator received SIGTERM; "
-                    "writing INTERRUPTED final snapshot"
-                )
-                loop.create_task(_signal_finalize())
-
-            async def _signal_finalize() -> None:
-                try:
-                    # Mirror the ENDED-driven path: refresh
-                    # tracked_duration_ns from the table BEFORE
-                    # publish_final, otherwise an interrupted run whose
-                    # STOP_PERFORMANCE_TRACKING never fired would
-                    # report duration_ns=0 and QPS=N/A in the final
-                    # report even after processing many tracked samples.
-                    registry.set_counter(
-                        MetricCounterKey.TRACKED_DURATION_NS.value,
-                        aggregator._table.total_tracked_duration_ns,
-                    )
-                    await publisher.publish_final(
-                        registry,
-                        n_pending_tasks=aggregator._table.in_flight_tasks_count,
-                        interrupted=True,
-                    )
-                except Exception:  # noqa: BLE001 — best-effort.
-                    logger.exception(
-                        "metrics aggregator: SIGTERM-triggered publish_final failed"
-                    )
-                shutdown_event.set()
-
-            loop.add_signal_handler(signal.SIGTERM, _on_sigterm)
+            on_sigterm, _sigterm_tasks = _make_sigterm_handler(
+                loop=loop,
+                registry=registry,
+                publisher=publisher,
+                table=aggregator._table,
+                shutdown_event=shutdown_event,
+            )
+            loop.add_signal_handler(signal.SIGTERM, on_sigterm)
             # No-op SIGINT handler: silence the default KeyboardInterrupt
             # and let the parent's ENDED-driven path drive shutdown.
             loop.add_signal_handler(

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

@@ -368,26 +368,18 @@ async def process(self, records: list[EventRecord]) -> None:
             # that fires before publish_final reflects the new state.
             self._session_state = SessionState.DRAINING
             logger.info("Draining %d async tasks...", table.in_flight_tasks_count)
-            try:
-                await asyncio.wait_for(
-                    table.drain_tasks(), timeout=self._drain_timeout_s
-                )
-            except TimeoutError:
+            # drain_tasks owns the timeout + cancel-and-await sequence so
+            # the pending count is captured BEFORE done-callbacks empty
+            # the in-flight set. Reading in_flight_tasks_count out here
+            # would always be 0 (see drain_tasks docstring).
+            n_pending = await table.drain_tasks(timeout=self._drain_timeout_s)
+            if n_pending > 0:
                 logger.warning(
-                    "drain_tasks timed out after %.1fs; some async metrics "
-                    "may be incomplete",
+                    "drain_tasks timed out after %.1fs; %d async tasks "
+                    "did not complete and were cancelled",
                     self._drain_timeout_s,
+                    n_pending,
                 )
-                # cancel() only *schedules* cancellation at the next await
-                # point. Await the cancelled tasks so they actually exit
-                # before publish_final reads n_pending — otherwise the
-                # snapshot reports stale-high pending counts and the
-                # event-loop tear-down emits "Task was destroyed but it
-                # is pending!" warnings on the cancelled set.
-                cancelled = table.cancel_in_flight_tasks()
-                if cancelled:
-                    await asyncio.gather(*cancelled, return_exceptions=True)
-            n_pending = table.in_flight_tasks_count
             logger.info(
                 "Async tasks drained (n_pending_tasks=%d at finalize)", n_pending
             )

src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py

@@ -494,26 +494,37 @@ def in_flight_tasks_count(self) -> int:
         """Number of async trigger tasks currently in flight."""
         return len(self._in_flight_tasks)
 
-    async def drain_tasks(self) -> None:
-        """Await all in-flight async trigger tasks."""
-        if self._in_flight_tasks:
+    async def drain_tasks(self, *, timeout: float | None = None) -> int:
+        """Await in-flight async trigger tasks.
+
+        With ``timeout``, the pending set at the timeout boundary is
+        cancelled and awaited; the count of those pending tasks is
+        returned (>0 indicates the drain timed out). Without
+        ``timeout``, blocks indefinitely and returns 0 on clean drain.
+
+        The pending count must be captured BEFORE the cancel-and-await
+        step: each task's ``add_done_callback(_in_flight_tasks.discard)``
+        empties ``_in_flight_tasks`` as cancellation propagates, so
+        reading ``in_flight_tasks_count`` after this method returns
+        would always be 0 — making a drain timeout indistinguishable
+        from a clean run.
+        """
+        if not self._in_flight_tasks:
+            return 0
+        if timeout is None:
             await asyncio.gather(*self._in_flight_tasks, return_exceptions=True)
             self._in_flight_tasks.clear()
-
-    def cancel_in_flight_tasks(self) -> list[asyncio.Task]:
-        """Cancel every in-flight async trigger task that hasn't finished.
-
-        Returns the tasks that were cancelled so callers can await them
-        (cancellation is only scheduled by ``Task.cancel()`` — the tasks
-        must still be awaited at a later point for the cancellation to
-        actually take effect).
-        """
-        cancelled: list[asyncio.Task] = []
-        for t in list(self._in_flight_tasks):
-            if not t.done():
+            return 0
+        _, still_pending = await asyncio.wait(
+            list(self._in_flight_tasks), timeout=timeout
+        )
+        n_pending = len(still_pending)
+        if still_pending:
+            for t in still_pending:
                 t.cancel()
-                cancelled.append(t)
-        return cancelled
+            await asyncio.gather(*still_pending, return_exceptions=True)
+        self._in_flight_tasks.clear()
+        return n_pending
 
     # --- Internal ---
 

src/inference_endpoint/metrics/report.py

@@ -296,23 +296,32 @@ def _display_metric(
     scale_factor: float = 1.0,
     newline: str = "",
 ) -> None:
+    # ``_scrub_nonfinite`` (snapshot.py) maps producer-side NaN/±Inf to
+    # ``None`` so the persisted JSON stays strict. Any of the named
+    # scalars / percentile values below can therefore be ``None`` —
+    # render an ``N/A`` indicator instead of crashing on
+    # ``None * scale_factor``.
+    def _scaled(v: Any) -> str:
+        if v is None:
+            return "N/A"
+        return f"{v * scale_factor:.2f}"
+
     for name, key in [
         ("Min", "min"),
         ("Max", "max"),
         ("Median", "median"),
         ("Avg.", "avg"),
         ("Std Dev.", "std_dev"),
     ]:
-        fn(f"  {name}: {metric_dict[key] * scale_factor:.2f} {unit}{newline}")
+        fn(f"  {name}: {_scaled(metric_dict[key])} {unit}{newline}")
 
     fn(f"\n  Histogram:{newline}")
     buckets = metric_dict["histogram"]["buckets"]
     counts = metric_dict["histogram"]["counts"]
 
     if buckets:
         bucket_strs = [
-            f"  [{lo * scale_factor:.2f}, {hi * scale_factor:.2f}"
-            + ("]" if i == len(buckets) - 1 else ")")
+            f"  [{_scaled(lo)}, {_scaled(hi)}" + ("]" if i == len(buckets) - 1 else ")")
             for i, (lo, hi) in enumerate(buckets)
         ]
         max_count = max(counts)
@@ -325,4 +334,4 @@ def _display_metric(
 
     fn(f"\n  Percentiles:{newline}")
     for p, val in metric_dict.get("percentiles", {}).items():
-        fn(f"  {p:>6}: {val * scale_factor:.2f} {unit}{newline}")
+        fn(f"  {p:>6}: {_scaled(val)} {unit}{newline}")

tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py

@@ -1055,3 +1055,68 @@ async def test_shutdown_drains_async_tasks(self, tmp_path):
     # exercised here. Adding a MockTokenizePool that raises on
     # token_count_async would let us assert no metric is emitted, the
     # aggregator does not crash, and the task set is cleaned up.
+
+    @pytest.mark.asyncio
+    async def test_drain_timeout_reports_pending_count(self, tmp_path):
+        """On drain timeout, publish_final must receive n_pending_tasks > 0.
+
+        AGENTS.md and the ``MetricsSnapshot.n_pending_tasks`` docstring
+        document the consumer contract: a drain-timeout run is detected
+        downstream as ``state == COMPLETE and n_pending_tasks > 0``. If
+        the producer always reports 0 here, the timeout is silently
+        rebadged as a clean run and the Report shows no warning.
+        """
+        loop = asyncio.get_event_loop()
+
+        class BlockingTokenizePool:
+            async def token_count_async(self, text, _loop):
+                await asyncio.sleep(10.0)  # exceeds drain timeout
+                return 0
+
+            def token_count(self, text):
+                return 0
+
+            def close(self):
+                pass
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *args):
+                self.close()
+
+        with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx:
+            agg, _, publisher = make_aggregator(
+                ctx,
+                loop,
+                "agg_drain_timeout",
+                tokenize_pool=BlockingTokenizePool(),
+            )
+            agg._drain_timeout_s = 0.05
+            try:
+                await agg.process(
+                    [
+                        session_event(
+                            SessionEventType.START_PERFORMANCE_TRACKING, ts=0
+                        ),
+                        sample_event(
+                            SampleEventType.ISSUED,
+                            "s1",
+                            ts=1000,
+                            data=PromptData(text="some text to tokenize"),
+                        ),
+                    ]
+                )
+                assert (
+                    agg._table.in_flight_tasks_count > 0
+                ), "precondition: ISL task must be in-flight before ENDED"
+                await agg.process([session_event(SessionEventType.ENDED, ts=2000)])
+
+                publisher.publish_final.assert_awaited_once()
+                kwargs = publisher.publish_final.await_args.kwargs
+                assert kwargs["n_pending_tasks"] > 0, (
+                    f"drain timeout must report stuck tasks; got "
+                    f"n_pending_tasks={kwargs['n_pending_tasks']}"
+                )
+            finally:
+                agg.close()