[Observability][API server] Stage latencies and per-modality output at /metrics

docs/serving.rst

@@ -533,6 +533,32 @@ warmup requests. Every family has a ``model_name`` label (the config's ``model``
   reaching the API server. It is measured on the server, so it does not include the
   network to the client, and a non-streaming client sees nothing until the end.
 
+Where a successful request's time went, from the timestamps every request already
+carries (same ``model_name`` and ``output_modalities`` labels; the conductor's
+timestamps are compared with the API server's, which assumes both run on one host, as
+they do today):
+
+- ``mstar_request_preprocess_seconds``: histogram, from submission to the data worker
+  handing the preprocessed request to the conductor, including its wait in the data
+  worker's queue.
+- ``mstar_request_engine_seconds``: histogram, from the conductor taking the request in
+  to the conductor finishing it.
+- ``mstar_request_delivery_seconds``: histogram, from the conductor finishing the
+  request to the result leaving the API server (the last chunks' postprocessing and
+  transfer, and the client handler picking it up).
+
+Output, per output modality (``modality`` is the chunk's modality, e.g. ``audio``):
+
+- ``mstar_output_chunks_total{model_name, modality}`` and
+  ``mstar_output_bytes_total{model_name, modality}``: counters of the result chunks (and
+  their bytes) that reached the API server, for every request whatever its outcome;
+  error chunks are not counted.
+- ``mstar_output_audio_seconds_total{model_name}``: counter, seconds of audio in those
+  chunks (16-bit PCM at the chunk's sample rate and channel count).
+- ``mstar_time_to_first_chunk_seconds{model_name, modality}``: histogram over
+  successful requests, from submission to the first chunk of each output modality
+  reaching the API server (time to first audio, first image, first video frame).
+
 .. code-block:: bash
 
    mstar serve kokoro --enable-metrics

mstar/api_server/data_worker.py

@@ -160,6 +160,8 @@ def __init__(
         self.discard_tensor_queue = queue.Queue()
         self.output_queue = queue.Queue()
         self.profile_queue = queue.Queue()
+        # Set by report_preprocess_finish(); read by the worker thread.
+        self.timing_requested = threading.Event()
         self.stop_event = threading.Event()
         # Wakes the thread's idle wait: every producer below signals it.
         self.wakeup = EventWakeup()
@@ -205,12 +207,18 @@ def __init__(
                 tensor_manager=self.tensor_manager,
                 model=model,
                 enable_prof=enable_prof,
+                timing_requested=self.timing_requested,
                 model_config=model_config,
                 enable_nvtx=enable_nvtx,
             )
         )
         self.thread.start()
 
+    def report_preprocess_finish(self) -> None:
+        """Put each request's preprocess-finish time on the profile queue
+        without ``--log-stats`` too (``mstar_request_preprocess_seconds``)."""
+        self.timing_requested.set()
+
     def _signal(self):
         """End the thread's idle wait: new work is queued. getattr: tests
         build this facade without its constructor."""
@@ -432,6 +440,7 @@ def __init__(
         model_config: dict | None = None,
         enable_nvtx: bool=False,
         wakeup: EventWakeup | None = None,
+        timing_requested: threading.Event | None = None,
     ):
         # keying a stream needs the deployment's page size as well as the
         # model's declaration, so a worker built without a config keys no stream
@@ -466,6 +475,8 @@ def __init__(
         self.device = device
         self.model = model
         self.enable_prof = enable_prof
+        # Preprocess-finish times are reported when enable_prof or this is set.
+        self.timing_requested = timing_requested if timing_requested is not None else threading.Event()
         # This thread turns each output tensor into client-ready bytes. At 720p
         # that is an 11 MiB SHM read plus a postprocess copy per engine step,
         # downstream of the last worker-side NVTX range.
@@ -644,11 +655,11 @@ def _process_input(
         # the *raw* inputs. ``perf_counter`` is consistent here because the worker
         # runs as a thread inside the API server process. (tx/rx are snapshotted
         # directly by the main thread at request completion — see APIServer.)
-        if self.enable_prof:
+        if self.enable_prof or self.timing_requested.is_set():
             self.profile_queue.put(DataWorkerProfile(
                 request_id=input.request_id,
                 preprocess_finish_time=time.perf_counter(),
-                inputs=self._summarize_inputs(input),
+                inputs=self._summarize_inputs(input) if self.enable_prof else [],
             ))
 
     def _prefix_keys(

mstar/api_server/entrypoint.py

@@ -237,6 +237,8 @@ class PendingRequest:
     # The client's request_id. The engine knows the request by its key, so
     # the message thread restamps each chunk with this before the client sees it.
     request_id: str | None = None
+    # Output modality -> perf_counter of its first chunk; kept only with metrics on.
+    first_chunk_times: dict[str, float] = field(default_factory=dict)
 
     def notify(self) -> None:
         """Wake the streaming handler, from any thread."""
@@ -497,6 +499,7 @@ def enable_metrics(self) -> None:
         self.metrics = ServingMetrics(
             self.model_name, in_flight=lambda: len(self.pending_requests),
         )
+        self.preprocess_worker.report_preprocess_finish()
 
     # ----------------------------------------------------------
     # Submitting a request
@@ -825,6 +828,7 @@ def _process_messages(self) -> None:
                 # the handlers are woken once per pass, after every chunk of
                 # the pass has been appended, not once per chunk
                 to_wake: dict[int, PendingRequest] = {}
+                metrics = self.metrics
                 for result_chunk in self.preprocess_worker.get_result_chunks():
                     logger.debug(
                         "Got result chunk of %s modality for request %s",
@@ -845,6 +849,8 @@ def _process_messages(self) -> None:
                         if req.profile.timing.first_chunk_time is None:
                             req.profile.timing.first_chunk_time = now
                         req.profile.timing.last_chunk_time = now
+                        if metrics is not None and result_chunk.modality != "error":
+                            req.first_chunk_times.setdefault(result_chunk.modality, now)
                         if req.request_id is not None:
                             result_chunk.request_id = req.request_id
                         if result_chunk.modality == "error":
@@ -870,6 +876,10 @@ def _process_messages(self) -> None:
                             # the data worker's per-request state once the
                             # client lets go of the request.
                             self.recently_completed[rid] = time.time()
+                    if metrics is not None and result_chunk.modality != "error":
+                        metrics.record_chunk(
+                            result_chunk.modality, result_chunk.data, result_chunk.metadata,
+                        )
                 if to_wake:
                     notify_many(to_wake.values())
             except Exception:

mstar/observability/metrics.py

@@ -18,7 +18,7 @@
 
 import time
 from collections.abc import Callable
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 try:
     from prometheus_client import (
@@ -59,6 +59,14 @@
     0.05, 0.1, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0, 5.0, 7.5, 10.0,
     15.0, 20.0, 30.0, 60.0, 120.0, 300.0, 600.0, 1200.0,
 )
+# Preprocess and delivery: under a millisecond for a short text prompt, up to
+# seconds for a long video upload or a large final payload.
+STAGE_BUCKETS = (
+    0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5,
+    1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0,
+)
+# 16-bit PCM: the data worker emits audio as headerless s16le.
+_AUDIO_BYTES_PER_SAMPLE = 2
 
 
 def require_prometheus_client() -> None:
@@ -90,6 +98,17 @@ def collect(self):
         yield family
 
 
+class _RequestSeries(NamedTuple):
+    """One output-modality set's labeled series, looked up once: a
+    ``labels()`` call costs about as much as the observation itself."""
+    outcomes: dict[str, Any]
+    e2e: Any
+    first_output: Any
+    preprocess: Any
+    engine: Any
+    delivery: Any
+
+
 class ServingMetrics:
     """The API server's registry and its request metrics."""
 
@@ -120,27 +139,117 @@ def __init__(self, model_name: str, in_flight: Callable[[], int]):
             registry=self.registry,
         )
         self.registry.register(_InFlight(model_name, in_flight))
-        self._seen_modalities: set[str] = set()
+
+        # Where a successful request's time went, from the timestamps it
+        # carries. The conductor's perf_counter is compared with this
+        # process's: both read the host's monotonic clock.
+        stages = ["model_name", "output_modalities"]
+        self.preprocess_latency = Histogram(
+            "mstar_request_preprocess_seconds",
+            "Successful requests: submission to the data worker handing the "
+            "request to the conductor, including its wait in the data worker's queue.",
+            stages, buckets=STAGE_BUCKETS, registry=self.registry,
+        )
+        self.engine_latency = Histogram(
+            "mstar_request_engine_seconds",
+            "Successful requests: conductor ingest to conductor finish.",
+            stages, buckets=E2E_BUCKETS, registry=self.registry,
+        )
+        self.delivery_latency = Histogram(
+            "mstar_request_delivery_seconds",
+            "Successful requests: conductor finish to the result leaving the API server.",
+            stages, buckets=STAGE_BUCKETS, registry=self.registry,
+        )
+
+        # Output volume, counted per chunk as it reaches the API server, for
+        # every request; modality names come from the model's graph.
+        self.output_chunks = Counter(
+            "mstar_output_chunks",
+            "Result chunks that reached the API server, by modality.",
+            ["model_name", "modality"], registry=self.registry,
+        )
+        self.output_bytes = Counter(
+            "mstar_output_bytes",
+            "Bytes of result chunks that reached the API server, by modality.",
+            ["model_name", "modality"], registry=self.registry,
+        )
+        self.output_audio_seconds = Counter(
+            "mstar_output_audio_seconds",
+            "Seconds of audio in the result chunks that reached the API server.",
+            ["model_name"], registry=self.registry,
+        ).labels(model_name)
+        self.time_to_first_chunk = Histogram(
+            "mstar_time_to_first_chunk_seconds",
+            "Successful requests: submission to the first result chunk of each "
+            "output modality reaching the API server.",
+            ["model_name", "modality"], buckets=FIRST_OUTPUT_BUCKETS, registry=self.registry,
+        )
+        # Label lookups resolved once per modality set or modality. Two
+        # threads may resolve the same key at once; labels() returns the same
+        # series to both, so the race is harmless.
+        self._request_series: dict[str, _RequestSeries] = {}
+        self._first_chunk_series: dict[str, Any] = {}
+        self._chunk_series: dict[str, tuple] = {}
+
+    def _series(self, modalities: str) -> _RequestSeries:
+        series = self._request_series.get(modalities)
+        if series is None:
+            labels = (self.model_name, modalities)
+            series = _RequestSeries(
+                # Exported at 0 from here on: rate() and increase() miss the
+                # first increment of a series that appears at 1.
+                outcomes={o: self.requests.labels(*labels, o) for o in OUTCOMES},
+                e2e=self.e2e_latency.labels(*labels),
+                first_output=self.time_to_first_output.labels(*labels),
+                preprocess=self.preprocess_latency.labels(*labels),
+                engine=self.engine_latency.labels(*labels),
+                delivery=self.delivery_latency.labels(*labels),
+            )
+            self._request_series[modalities] = series
+        return series
 
     def record(self, req: PendingRequest, outcome: str) -> None:
         """Count one request that left the API server's table."""
         now = time.perf_counter()
-        modalities = modalities_label(req.output_modalities)
-        if modalities not in self._seen_modalities:
-            # rate() and increase() miss the first increment of a series that
-            # appears at 1, so every outcome is exported at 0 from here on.
-            for o in OUTCOMES:
-                self.requests.labels(self.model_name, modalities, o)
-            self._seen_modalities.add(modalities)
-        self.requests.labels(self.model_name, modalities, outcome).inc()
+        series = self._series(modalities_label(req.output_modalities))
+        series.outcomes[outcome].inc()
         if outcome != "success":
             return
         timing = req.profile.timing
-        self.e2e_latency.labels(self.model_name, modalities).observe(
-            now - timing.recv_time)
+        series.e2e.observe(now - timing.recv_time)
         if timing.first_chunk_time is not None:
-            self.time_to_first_output.labels(self.model_name, modalities).observe(
-                timing.first_chunk_time - timing.recv_time)
+            series.first_output.observe(timing.first_chunk_time - timing.recv_time)
+        if timing.preprocess_finish_time is not None:
+            series.preprocess.observe(timing.preprocess_finish_time - timing.recv_time)
+        # Every successful request got its RequestComplete, which carries
+        # both; a value from another host's clock would come out negative.
+        if timing.conductor_finish_time is not None and now >= timing.conductor_finish_time:
+            series.engine.observe(timing.conductor_finish_time - timing.conductor_ingest_time)
+            series.delivery.observe(now - timing.conductor_finish_time)
+        for modality, first in req.first_chunk_times.items():
+            first_chunk = self._first_chunk_series.get(modality)
+            if first_chunk is None:
+                first_chunk = self.time_to_first_chunk.labels(self.model_name, modality)
+                self._first_chunk_series[modality] = first_chunk
+            first_chunk.observe(first - timing.recv_time)
+
+    def record_chunk(self, modality: str, data: bytes, metadata: dict | None) -> None:
+        """Count one result chunk that reached the API server (not an error chunk)."""
+        series = self._chunk_series.get(modality)
+        if series is None:
+            series = (self.output_chunks.labels(self.model_name, modality),
+                      self.output_bytes.labels(self.model_name, modality))
+            self._chunk_series[modality] = series
+        chunks, nbytes = series
+        chunks.inc()
+        nbytes.inc(len(data))
+        if modality == "audio" and metadata:
+            # Chunks from the result-token path carry no sample rate.
+            rate = metadata.get("sample_rate")
+            if rate:
+                channels = metadata.get("num_channels") or 1
+                self.output_audio_seconds.inc(
+                    len(data) / (_AUDIO_BYTES_PER_SAMPLE * channels * rate))
 
     def render(self) -> bytes:
         return generate_latest(self.registry)

test/modular/test_delivery_chunk_counts.py

@@ -302,6 +302,7 @@ def test_completion_that_overtakes_the_last_chunks_waits_for_them():
     server.recently_completed = collections.OrderedDict()
     server._recently_completed_ttl = 3600.0
     server.log_stats = False
+    server.metrics = None
     server.enable_nvtx = False
     server.preprocess_worker = worker
     req = PendingRequest(streaming=False, input_modalities=["text", "image", "text"],