[Observability][Data worker] Preprocess time, failures, upload sizes and queue depth at /metrics

docs/serving.rst

@@ -592,6 +592,38 @@ Output, per output modality (``modality`` is the chunk's modality, e.g. ``audio`
   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).
 
+The data worker (a thread in the API server that loads each request's media and runs
+the model's prompt processing, preprocessing one request at a time) reports its own
+work.
+``input_modalities`` is the request's input modalities, sorted and comma-joined like
+``output_modalities``:
+
+- ``mstar_data_worker_preprocess_seconds{model_name, input_modalities}``: histogram over
+  requests whose preprocessing finished, from the data worker taking the request off its
+  queue to handing it to the conductor. Unlike ``mstar_request_preprocess_seconds`` it
+  leaves out the wait in the queue, and it counts requests whatever their outcome later.
+- ``mstar_data_worker_preprocess_failures_total{model_name, input_modalities}``:
+  counter, requests whose preprocessing failed in the data worker, each once, whatever
+  raised: a file count or media length over its limit, a file that cannot be decoded,
+  the model's prompt processing, or handing the inputs on. The client gets the error
+  (a 400 for a bad input), so the request also ends as ``outcome="error"`` in
+  ``mstar_requests_total`` unless its client had already gone. That counter counts
+  failures from every stage by output modalities; this one only the data worker's
+  preprocessing, by input modalities (failures turning results into output chunks are
+  not counted here).
+- ``mstar_upload_size_bytes{model_name, modality}``: histogram, the size on disk of
+  each image, audio or video file a request carried to the data worker (``modality`` is
+  ``image``, ``audio`` or ``video``), whether the request succeeded or failed. Every
+  route writes its media to a file first: an upload as received, inline media (a
+  ``data:`` URL, base64 files on ``/generate/ws``) after base64 decoding. The realtime
+  route and long transcriptions send WAV files they write from the decoded audio, so
+  those count the WAV sizes. Requests refused before they were submitted (an oversized
+  body, a bad form) are not counted.
+- ``mstar_data_worker_queue_depth{model_name}``: gauge, read at scrape time, requests
+  submitted and not yet taken up by the data worker (the one it is working on is not
+  counted). A request whose client went away still counts until the data worker takes
+  it up.
+
 .. code-block:: bash
 
    mstar serve kokoro --enable-metrics

mstar/api_server/data_worker.py

@@ -216,9 +216,15 @@ def __init__(
 
     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``)."""
+        without ``--log-stats`` too (``mstar_request_preprocess_seconds``),
+        with the data worker's own time, a failed preprocess, and the input
+        files' sizes."""
         self.timing_requested.set()
 
+    def queue_depth(self) -> int:
+        """Requests submitted and not yet taken up by the data worker thread."""
+        return self.request_input_queue.qsize()
+
     def _signal(self):
         """End the thread's idle wait: new work is queued. getattr: tests
         build this facade without its constructor."""
@@ -475,7 +481,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.
+        # Preprocess-finish times are reported when enable_prof or this is set,
+        # the data worker's own metrics only when 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,
@@ -529,7 +536,7 @@ def _drop_request_state(self, request_id: str) -> None:
                 state.pop(request_id, None)
 
     def _process_input(
-        self, input: PreprocessInput
+        self, input: PreprocessInput, profile: DataWorkerProfile | None = None,
     ):
         tensors: NameToTensorList = {}
         input_metadata = {}
@@ -655,12 +662,38 @@ 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 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) if self.enable_prof else [],
-            ))
+        # ``profile`` was started at dequeue when metrics are on.
+        if profile is None and self.enable_prof:
+            profile = DataWorkerProfile(request_id=input.request_id)
+        if profile is not None:
+            profile.preprocess_finish_time = time.perf_counter()
+            if self.enable_prof:
+                profile.inputs = self._summarize_inputs(input)
+            self.profile_queue.put(profile)
+
+    @staticmethod
+    def _start_profile(input: PreprocessInput) -> DataWorkerProfile:
+        """The metrics side of a request's profile, taken as it leaves the
+        queue: the HTTP layer deletes uploads on a timer, so their sizes are
+        read before the work, not after."""
+        dequeue_time = time.perf_counter()
+        sizes: dict[str, list[int]] = {}
+        for modality, paths in (input.file_paths or {}).items():
+            # The modalities _process_input loads; it skips any other key.
+            if modality not in ("image", "audio", "video"):
+                continue
+            found = sizes[modality] = []
+            for path in paths:
+                try:
+                    found.append(os.stat(path).st_size)
+                except (OSError, ValueError):
+                    pass  # gone or not a valid path: its load fails, and the failure is counted
+        return DataWorkerProfile(
+            request_id=input.request_id,
+            dequeue_time=dequeue_time,
+            input_modalities=input.input_modalities,
+            upload_sizes=sizes,
+        )
 
     def _prefix_keys(
         self, tensors: dict, layouts: dict, file_paths: dict | None,
@@ -1272,9 +1305,15 @@ def run(self):
                     # Releases queued before this request was submitted belong
                     # to an earlier request under its id: apply them first.
                     self._process_releases()
+                    profile = None
                     try:
-                        self._process_input(pre_input)
+                        if self.timing_requested.is_set():
+                            profile = self._start_profile(pre_input)
+                        self._process_input(pre_input, profile)
                     except Exception as exc:  # noqa: BLE001 — any failure must reach the client
+                        if profile is not None:
+                            profile.failed = True
+                            self.profile_queue.put(profile)
                         # A request whose media load or prompt processing fails
                         # never reaches the conductor, so nothing downstream
                         # would ever complete it; surface the failure as an

mstar/api_server/entrypoint.py

@@ -508,6 +508,7 @@ def enable_metrics(self) -> None:
         compile times stay out of the latency histograms."""
         self.metrics = ServingMetrics(
             self.model_name, in_flight=lambda: len(self.pending_requests),
+            data_worker_queue_depth=self.preprocess_worker.queue_depth,
         )
         self.preprocess_worker.report_preprocess_finish()
 
@@ -839,7 +840,14 @@ def _process_messages(self) -> None:
                 # Apply data-worker profiling: preprocess-finish timestamp + raw
                 # input sizes. (The data worker's own tx/rx are snapshotted
                 # directly in _prune_recently_completed once the request is done.)
+                metrics = self.metrics
                 for update in self.preprocess_worker.get_profile_updates():
+                    # A --log-stats profile begun before metrics were on has
+                    # no dequeue time.
+                    if metrics is not None and update.dequeue_time is not None:
+                        metrics.record_preprocess(update)
+                    if update.failed:
+                        continue
                     with self.request_lock:
                         req = self.pending_requests.get(update.request_id)
                         if req is not None:
@@ -850,7 +858,6 @@ 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",

mstar/api_server/request_types.py

@@ -97,13 +97,21 @@ class APIServerMessage:
 
 @dataclass
 class DataWorkerProfile:
-    """Profiling reported by the API-server data worker at preprocess finish:
-    the timestamp at which the request was handed to the conductor and the
-    per-modality sizes of the raw inputs. (The data worker's tx/rx are read
-    directly from its tensor manager at request completion, not via this.)"""
+    """Profiling reported by the API-server data worker at preprocess finish
+    (or failure): the timestamp at which the request was handed to the
+    conductor and the per-modality sizes of the raw inputs. (The data worker's
+    tx/rx are read directly from its tensor manager at request completion,
+    not via this.)"""
     request_id: str
     preprocess_finish_time: float | None = None  # time.perf_counter
     inputs: list[InputInfo] = field(default_factory=list)
+    # Set only with metrics on: the data worker's own numbers for the request.
+    dequeue_time: float | None = None  # perf_counter, taken off the data worker's queue
+    input_modalities: list[str] = field(default_factory=list)
+    # modality -> size in bytes of each input file, read when it was dequeued
+    upload_sizes: dict[str, list[int]] = field(default_factory=dict)
+    # preprocessing raised: the client gets an error chunk instead
+    failed: bool = False
 
 
 @dataclass

mstar/observability/metrics.py

@@ -56,7 +56,7 @@
 
 if TYPE_CHECKING:
     from mstar.api_server.entrypoint import PendingRequest
-    from mstar.api_server.request_types import ConductorStatus
+    from mstar.api_server.request_types import ConductorStatus, DataWorkerProfile
 
 logger = logging.getLogger(__name__)
 
@@ -84,6 +84,14 @@
     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,
 )
+# Input files: a few KiB for a small image to MSTAR_MAX_BODY_MB (128 MiB by
+# default) and the transcription route's larger cap.
+_MiB = 1 << 20
+UPLOAD_SIZE_BUCKETS = (
+    1 << 10, 4 << 10, 16 << 10, 32 << 10, 64 << 10, 128 << 10, 256 << 10, 512 << 10,
+    _MiB, 2 * _MiB, 4 * _MiB, 8 * _MiB, 16 * _MiB, 32 * _MiB, 64 * _MiB, 128 * _MiB,
+    256 * _MiB, 1024 * _MiB,
+)
 # 16-bit PCM: the data worker emits audio as headerless s16le.
 _AUDIO_BYTES_PER_SAMPLE = 2
 
@@ -94,25 +102,22 @@ def require_prometheus_client() -> None:
 
 
 def modalities_label(modalities: list[str]) -> str:
-    """One label value per set of output modalities, e.g. ``audio,text``."""
+    """One label value per set of input or output modalities, e.g. ``audio,text``."""
     return ",".join(sorted(set(modalities)))
 
 
-class _InFlight:
-    """``mstar_requests_in_flight``, read when scraped rather than counted up
-    and down, so a request that ends on an unusual path cannot leave it off."""
+class _ScrapedGauge:
+    """A gauge read when scraped rather than counted up and down, so a request
+    that ends on an unusual path cannot leave it off."""
 
-    def __init__(self, model_name: str, count: Callable[[], int]):
+    def __init__(self, name: str, documentation: str, model_name: str, count: Callable[[], int]):
+        self.name = name
+        self.documentation = documentation
         self.model_name = model_name
         self.count = count
 
     def collect(self):
-        family = GaugeMetricFamily(
-            "mstar_requests_in_flight",
-            "Requests the API server holds: submitted and not yet returned, "
-            "failed or aborted.",
-            labels=["model_name"],
-        )
+        family = GaugeMetricFamily(self.name, self.documentation, labels=["model_name"])
         family.add_metric([self.model_name], self.count())
         yield family
 
@@ -338,10 +343,19 @@ async def send_and_note_status(message):
                 metrics.refusals.inc(route_template(scope, self.routes()), status)
 
 
+class _DataWorkerSeries(NamedTuple):
+    """One input-modality set's data worker series."""
+    seconds: Any
+    failures: Any
+
+
 class ServingMetrics:
     """The API server's registry and its request metrics."""
 
-    def __init__(self, model_name: str, in_flight: Callable[[], int]):
+    def __init__(
+        self, model_name: str, in_flight: Callable[[], int],
+        data_worker_queue_depth: Callable[[], int],
+    ):
         require_prometheus_client()
         self.model_name = model_name
         self.content_type = CONTENT_TYPE_LATEST
@@ -372,7 +386,12 @@ def __init__(self, model_name: str, in_flight: Callable[[], int]):
             queue.name, queue.help, ["model_name", *queue.labels],
             buckets=queue.buckets, registry=self.registry,
         )
-        self.registry.register(_InFlight(model_name, in_flight))
+        self.registry.register(_ScrapedGauge(
+            "mstar_requests_in_flight",
+            "Requests the API server holds: submitted and not yet returned, "
+            "failed or aborted.",
+            model_name, in_flight,
+        ))
         self.refusals = _Refusals(model_name)
         self.registry.register(self.refusals)
         self._conductor = _ConductorStatus(model_name)
@@ -422,12 +441,41 @@ def __init__(self, model_name: str, in_flight: Callable[[], int]):
             "output modality reaching the API server.",
             ["model_name", "modality"], buckets=FIRST_OUTPUT_BUCKETS, registry=self.registry,
         )
+        # The data worker's own work, from the profile it reports per request
+        # it took off its queue; perf_counter is shared, it is a thread here.
+        data_worker = ["model_name", "input_modalities"]
+        self.data_worker_seconds = Histogram(
+            "mstar_data_worker_preprocess_seconds",
+            "Requests the data worker preprocessed: from taking the request off "
+            "its queue to handing it to the conductor.",
+            data_worker, buckets=STAGE_BUCKETS, registry=self.registry,
+        )
+        self.data_worker_failures = Counter(
+            "mstar_data_worker_preprocess_failures",
+            "Requests whose preprocessing failed in the data worker (media "
+            "limits, undecodable media, the model's prompt processing, any other error).",
+            data_worker, registry=self.registry,
+        )
+        self.upload_size = Histogram(
+            "mstar_upload_size_bytes",
+            "Size of each image, audio or video file a request carried to the "
+            "data worker, uploaded or inline.",
+            ["model_name", "modality"], buckets=UPLOAD_SIZE_BUCKETS, registry=self.registry,
+        )
+        self.registry.register(_ScrapedGauge(
+            "mstar_data_worker_queue_depth",
+            "Requests submitted and waiting for the data worker to take them up.",
+            model_name, data_worker_queue_depth,
+        ))
+
         # 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] = {}
+        self._data_worker_series: dict[str, _DataWorkerSeries] = {}
+        self._upload_series: dict[str, Any] = {}
 
     def _series(self, modalities: str) -> _RequestSeries:
         series = self._request_series.get(modalities)
@@ -498,5 +546,32 @@ def record_chunk(self, modality: str, data: bytes, metadata: dict | None) -> Non
                 self.output_audio_seconds.inc(
                     len(data) / (_AUDIO_BYTES_PER_SAMPLE * channels * rate))
 
+    def record_preprocess(self, profile: DataWorkerProfile) -> None:
+        """Record the data worker's numbers for one request it took off its
+        queue: its time if preprocessing finished, else a failure, and the size
+        of each input file either way."""
+        modalities = modalities_label(profile.input_modalities)
+        series = self._data_worker_series.get(modalities)
+        if series is None:
+            labels = (self.model_name, modalities)
+            # Both from the first request of a modality set, so the first
+            # failure is not lost to rate().
+            series = _DataWorkerSeries(
+                self.data_worker_seconds.labels(*labels),
+                self.data_worker_failures.labels(*labels),
+            )
+            self._data_worker_series[modalities] = series
+        if profile.failed:
+            series.failures.inc()
+        else:
+            series.seconds.observe(profile.preprocess_finish_time - profile.dequeue_time)
+        for modality, sizes in profile.upload_sizes.items():
+            upload = self._upload_series.get(modality)
+            if upload is None:
+                upload = self.upload_size.labels(self.model_name, modality)
+                self._upload_series[modality] = upload
+            for size in sizes:
+                upload.observe(size)
+
     def render(self) -> bytes:
         return generate_latest(self.registry)

test/modular/test_api_completion_guard.py

@@ -52,6 +52,7 @@ def process_prompt(self, *args, **kwargs):
     wt.reads_done_queue = queue.Queue()
     wt.discard_tensor_queue = queue.Queue()
     wt.stop_event = threading.Event()
+    wt.timing_requested = threading.Event()
     wt.communicator = SimpleNamespace(
         get_all_new_messages=lambda: [], send=lambda *args: None,
     )