From 3f6befdf2e5c8d9e21de38aa2a33d572bd60868c Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Tue, 20 Oct 2026 08:03:53 +0000 Subject: [PATCH] [Observability][Data worker] Preprocess time, failures, upload sizes and queue depth at /metrics The data worker reports its own work per request when metrics are on: the profile it already puts on its profile queue now starts when the request leaves the data worker's queue, carries the request's input modalities and the size of each image, audio or video file (read at dequeue, before the HTTP layer's cleanup timer can remove them), and is also put, marked failed, when preprocessing raises. The API server's message thread, which already drains that queue, records the data worker's time from dequeue to the conductor handoff, counts failures by input modality set and observes each file's size. The data worker queue depth is read from its input queue at scrape time. With metrics off the data worker does nothing new beyond the existing flag check. --- docs/serving.rst | 32 +++ mstar/api_server/data_worker.py | 59 +++- mstar/api_server/entrypoint.py | 9 +- mstar/api_server/request_types.py | 16 +- mstar/observability/metrics.py | 103 ++++++- test/modular/test_api_completion_guard.py | 1 + test/modular/test_serving_metrics.py | 324 ++++++++++++++++++++-- test/rust/test_rust_frontend.py | 6 +- 8 files changed, 499 insertions(+), 51 deletions(-) diff --git a/docs/serving.rst b/docs/serving.rst index 3c6e3f5522..934f9c3352 100644 --- a/docs/serving.rst +++ b/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 diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index e198999b7b..b985bf9363 100644 --- a/mstar/api_server/data_worker.py +++ b/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 diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index fb026bead3..7ddf620b9c 100644 --- a/mstar/api_server/entrypoint.py +++ b/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", diff --git a/mstar/api_server/request_types.py b/mstar/api_server/request_types.py index 82a0209ce6..54521100eb 100644 --- a/mstar/api_server/request_types.py +++ b/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 diff --git a/mstar/observability/metrics.py b/mstar/observability/metrics.py index f16887ef4c..8063630716 100644 --- a/mstar/observability/metrics.py +++ b/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) diff --git a/test/modular/test_api_completion_guard.py b/test/modular/test_api_completion_guard.py index 8cba57be8f..e16321c007 100644 --- a/test/modular/test_api_completion_guard.py +++ b/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, ) diff --git a/test/modular/test_serving_metrics.py b/test/modular/test_serving_metrics.py index 6d552209bb..ae2ef88899 100644 --- a/test/modular/test_serving_metrics.py +++ b/test/modular/test_serving_metrics.py @@ -1,8 +1,9 @@ """Prometheus metrics at ``/metrics`` (``--enable-metrics``). Requests go through the FastAPI app and the API server's real message thread; -only the data worker and the conductor are stand-ins, fed the messages the -real ones send. Each test scrapes ``/metrics`` and reads the samples back. +the conductor is a stand-in fed the messages the real one sends, and so is the +data worker except in the data worker's own tests, which run its real thread. +Each test scrapes ``/metrics`` and reads the samples back. """ from __future__ import annotations @@ -10,6 +11,7 @@ import asyncio import contextlib import importlib +import os import queue import sys import threading @@ -54,6 +56,9 @@ def __init__(self): def report_preprocess_finish(self): self.reporting = True + def queue_depth(self): + return 0 + def new_request(self, inp): self.inputs.append(inp) if self.reporting: @@ -659,6 +664,26 @@ def test_with_metrics_off_no_chunk_is_timed(tmp_path): thread.join(timeout=5) +class _StubTensorManager: + """The data worker's transport, holding nothing: these tests stop at the + handoff to the conductor.""" + + def store_and_return_tensor_info(self, request_id, tensors): + return {} + + def register_for_send(self, request_id, tensor_infos): + pass + + def get_ready_tensors(self): + return {} + + def has_inflight_reads(self, request_id): + return False + + def force_cleanup_request(self, request_id): + pass + + def test_the_data_worker_reports_preprocess_finish_once_asked(tmp_path): """The real data worker thread stamps the handoff without ``--log-stats`` once ``report_preprocess_finish`` is called, and skips the input summary.""" @@ -669,22 +694,6 @@ class _Model: def process_prompt(self, *args, **kwargs): return None - class _TensorManager: - def store_and_return_tensor_info(self, request_id, tensors): - return {} - - def register_for_send(self, request_id, tensor_infos): - pass - - def get_ready_tensors(self): - return {} - - def has_inflight_reads(self, request_id): - return False - - def force_cleanup_request(self, request_id): - pass - sent = [] facade = PreprocessWorker.__new__(PreprocessWorker) facade.timing_requested = threading.Event() @@ -697,7 +706,7 @@ def force_cleanup_request(self, request_id): discard_tensor_queue=queue.Queue(), stop_event=stop, communicator=types.SimpleNamespace(get_all_new_messages=lambda: [], send=lambda entity, msg: sent.append(msg)), - tensor_manager=_TensorManager(), model=_Model(), + tensor_manager=_StubTensorManager(), model=_Model(), timing_requested=facade.timing_requested, ) thread = threading.Thread(target=worker.run, daemon=True) @@ -880,6 +889,281 @@ def test_the_rust_bridge_serves_the_conductor_families(served): assert _without_age(bridged) == _without_age(python_route) +SLOW_PROMPT_S = 0.03 + + +class _MediaModel: + """Decodes images with the base model's real loader and WAVs with the + stdlib (CI has neither soundfile nor torchcodec); takes one image per + request. The prompt text can slow, hold or fail its prompt processing.""" + + def __init__(self): + self.holding = threading.Event() + self.release = threading.Event() + + def max_input_files(self): + return {"image": 1} + + def load_image(self, filepath, device): + from mstar.model.base import Model + + return Model.load_image(self, filepath, device) + + def load_audio(self, filepath, device): + import wave + + import torch + + from mstar.model.base import TensorAndMetadata + + with wave.open(filepath) as f: + rate, pcm = f.getframerate(), f.readframes(f.getnframes()) + audio = torch.frombuffer(bytearray(pcm), dtype=torch.int16).float() / 32768 + return TensorAndMetadata(audio, metadata={"sample_rate": rate, "num_channels": 1}) + + def process_prompt(self, text, *args, **kwargs): + if text == "slow": + time.sleep(SLOW_PROMPT_S) + elif text == "hold": + self.holding.set() + assert self.release.wait(5) + elif text == "raise": + raise ValueError("the prompt is too long") + + +class _RealDataWorker: + """The API server on the real data worker: its facade and its thread, a + stand-in transport and conductor, and the API server's message thread.""" + + def __init__(self, tmp_path, *, metrics, enable_prof=False): + from mstar.api_server.data_worker import PreprocessWorker, PreprocessWorkerThread + + self.model, self.sent, self.tmp_path = _MediaModel(), [], tmp_path + facade = self.facade = PreprocessWorker.__new__(PreprocessWorker) + for name in ("request_input_queue", "result_tensor_input_queue", "cleanup_request_queue", + "abort_request_queue", "reads_done_queue", "discard_tensor_queue", + "output_queue", "profile_queue"): + setattr(facade, name, queue.Queue()) + facade.timing_requested, facade.stop_event = threading.Event(), threading.Event() + facade.per_request_reading_tensors, facade.received_chunk_counts = {}, {} + worker = PreprocessWorkerThread( + in_queue=facade.request_input_queue, result_tensor_queue=facade.result_tensor_input_queue, + out_queue=facade.output_queue, profile_queue=facade.profile_queue, + cleanup_request_queue=facade.cleanup_request_queue, + abort_request_queue=facade.abort_request_queue, + reads_done_queue=facade.reads_done_queue, + discard_tensor_queue=facade.discard_tensor_queue, stop_event=facade.stop_event, + communicator=types.SimpleNamespace(get_all_new_messages=lambda: [], + send=lambda entity, msg: self.sent.append(msg)), + tensor_manager=_StubTensorManager(), model=self.model, enable_prof=enable_prof, + timing_requested=facade.timing_requested, + ) + facade.thread = threading.Thread(target=worker.run, daemon=True) + facade.thread.start() + self.server = make_test_api_server( + preprocess_worker=facade, communicator=_Conductor(), upload_dir=tmp_path, + ) + if metrics: + self.server.enable_metrics() + self.message_thread = threading.Thread(target=self.server._process_messages, daemon=True) + self.message_thread.start() + + def submit(self, text, files=(), modalities=("text",)): + file_paths = {} + for modality, path in files: + file_paths.setdefault(modality, []).append(path) + return self.server.submit_request( + text=text, file_paths=file_paths or None, input_modalities=list(modalities), + output_modalities=["text"], streaming=False, + ) + + def scrape(self): + return _parse(self.server.metrics.render().decode()) + + def stop(self): + self.model.release.set() + self.server.running = False + self.facade.shutdown() + self.message_thread.join(timeout=5) + + +@pytest.fixture +def real_data_worker(monkeypatch, tmp_path): + pytest.importorskip("prometheus_client") + started = [] + + def start(**kwargs): + served = _RealDataWorker(tmp_path, **kwargs) + started.append(served) + monkeypatch.setattr(entrypoint, "api_server", served.server) + return served + + yield start + for served in started: + served.stop() + + +def _png(path, side=48): + import numpy as np + from PIL import Image + + pixels = np.random.default_rng(0).integers(0, 256, (side, side, 3), dtype=np.uint8) + Image.fromarray(pixels).save(path) + return str(path) + + +def _wav(path, seconds, rate=16_000): + import wave + + with wave.open(str(path), "wb") as out: + out.setnchannels(1) + out.setsampwidth(2) + out.setframerate(rate) + out.writeframes(b"\x01\x00" * int(seconds * rate)) + return str(path) + + +def _data_worker_hist(samples, modalities): + return _hist(samples, "mstar_data_worker_preprocess_seconds", input_modalities=modalities) + + +def _failures(samples, modalities): + return samples.get(("mstar_data_worker_preprocess_failures_total", + _labels(input_modalities=modalities)), 0) + + +def _upload_buckets(samples, modality): + return { + float(dict(k[1])["le"]): v for k, v in samples.items() + if k[0] == "mstar_upload_size_bytes_bucket" and dict(k[1])["modality"] == modality + } + + +def _preprocessed(samples): + """Requests the data worker finished or failed, by its metrics.""" + return sum(v for (name, _), v in samples.items() + if name in ("mstar_data_worker_preprocess_seconds_count", + "mstar_data_worker_preprocess_failures_total")) + + +def test_the_data_worker_times_preprocessing_and_counts_failures_by_input_modalities( + real_data_worker, monkeypatch, tmp_path): + """Each request the real data worker takes up is timed if it reaches the + conductor, or counted once as a failure (a corrupt upload through + ``/generate``, too many images, audio over the length limit, the prompt + processing raising); every image and audio file is sized either way.""" + monkeypatch.setenv("MSTAR_MAX_AUDIO_SECONDS", "2") + served = real_data_worker(metrics=True) + png, png2 = _png(tmp_path / "a.png"), _png(tmp_path / "b.png") + short, long = _wav(tmp_path / "short.wav", 1.0), _wav(tmp_path / "long.wav", 3.0) + corrupt = b"\x89PNG\r\n\x1a\n" + b"\x00" * 3000 + start = time.perf_counter() + + served.submit("slow") + served.submit("describe", [("image", png)], ("image", "text")) + served.submit("transcribe", [("audio", short)], ("audio", "text")) + r = TestClient(entrypoint.app).post( + "/generate", data={"text": "describe", "streaming": "false"}, + files={"files": ("upload.png", corrupt, "image/png")}, + ) + assert r.status_code == 400 and "image is not in a format" in r.text, r.text + served.submit("compare", [("image", png), ("image", png2)], ("image", "text")) + served.submit("transcribe", [("audio", long)], ("audio", "text")) + served.submit("raise") + _wait(lambda: _preprocessed(served.scrape()) == 7) + elapsed = time.perf_counter() - start + s = served.scrape() + + assert len(served.sent) == 3 # the three that reached the conductor + text_n, text_s = _data_worker_hist(s, "text") + assert text_n == 1 and SLOW_PROMPT_S <= text_s < elapsed + for modalities in ("image,text", "audio,text"): + n, seconds = _data_worker_hist(s, modalities) + assert n == 1 and 0 < seconds < elapsed, modalities + assert {m: _failures(s, m) for m in ("text", "image,text", "audio,text")} == { + "text": 1, "image,text": 2, "audio,text": 1} + + sizes = { + "image": [os.path.getsize(png)] * 2 + [os.path.getsize(png2), len(corrupt)], + "audio": [os.path.getsize(short), os.path.getsize(long)], + } + for modality, observed in sizes.items(): + assert _hist(s, "mstar_upload_size_bytes", modality=modality) == ( + len(observed), sum(observed)), modality + buckets = _upload_buckets(s, modality) + assert len(buckets) > 2 + assert buckets == {le: sum(size <= le for size in observed) for le in buckets}, modality + assert not _upload_buckets(s, "text") + + +def test_queue_depth_counts_requests_waiting_for_the_data_worker(real_data_worker, tmp_path): + """Read at scrape time: requests queued behind the one the data worker is + working on. Its preprocess time leaves their wait in the queue out.""" + served = real_data_worker(metrics=True) + clip = _wav(tmp_path / "clip.wav", 0.5) + assert served.scrape()[("mstar_data_worker_queue_depth", _labels())] == 0 + + served.submit("hold") + assert served.model.holding.wait(5) + queued_at = time.perf_counter() + served.submit("transcribe", [("audio", clip)], ("audio", "text")) + served.submit("transcribe", [("audio", clip)], ("audio", "text")) + assert served.scrape()[("mstar_data_worker_queue_depth", _labels())] == 2 + time.sleep(0.2) + waited = time.perf_counter() - queued_at + served.model.release.set() + _wait(lambda: _preprocessed(served.scrape()) == 3) + + s = served.scrape() + assert s[("mstar_data_worker_queue_depth", _labels())] == 0 + held_n, held_s = _data_worker_hist(s, "text") + assert held_n == 1 and held_s >= waited + queued_n, queued_s = _data_worker_hist(s, "audio,text") + assert queued_n == 2 and 0 < queued_s < waited + + +@pytest.mark.parametrize("enable_prof", [False, True], ids=["plain", "log_stats"]) +def test_with_metrics_off_the_data_worker_reports_nothing_new(real_data_worker, tmp_path, enable_prof): + """No profile for a failed preprocess and no metrics fields on a + successful one; ``--log-stats`` still gets its input summary.""" + served = real_data_worker(metrics=False, enable_prof=enable_prof) + png = _png(tmp_path / "a.png") + served.facade.get_profile_updates = lambda: [] # keep the updates for the test to read + key_ok = served.server._live[served.submit("describe", [("image", png)], ("image", "text"))] + failing = served.server.pending_requests[served.server._live[served.submit("raise")]] + # the failure's error chunk; the data worker finished the first request before + _wait(lambda: failing.error is not None) + assert len(served.sent) == 1 + + updates = [] + while not served.facade.profile_queue.empty(): + updates.append(served.facade.profile_queue.get()) + assert served.server.metrics is None + if not enable_prof: + assert updates == [] + return + [update] = updates + assert update.request_id == key_ok and update.preprocess_finish_time is not None + assert [(i.modality, i.count, i.total_bytes) for i in update.inputs] == [ + ("text", 1, len("describe")), ("image", 1, os.path.getsize(png))] + assert (update.dequeue_time, update.input_modalities, update.upload_sizes, update.failed) == ( + None, [], {}, False) + + +def test_log_stats_and_metrics_together_keep_the_input_summary(real_data_worker, tmp_path): + served = real_data_worker(metrics=True, enable_prof=True) + png = _png(tmp_path / "a.png") + key = served.server._live[served.submit("describe", [("image", png)], ("image", "text"))] + _wait(lambda: served.server.pending_requests[key].profile.inputs) + + inputs = served.server.pending_requests[key].profile.inputs + assert [(i.modality, i.total_bytes) for i in inputs] == [ + ("text", len("describe")), ("image", os.path.getsize(png))] + s = served.scrape() + assert _data_worker_hist(s, "image,text")[0] == 1 + assert _hist(s, "mstar_upload_size_bytes", modality="image") == (1, os.path.getsize(png)) + + class _StopError(Exception): pass @@ -1059,7 +1343,7 @@ def test_a_missing_package_is_reported_and_a_broken_one_raises(monkeypatch): importlib.reload(metrics_module) assert metrics_module.MISSING_DEPENDENCY is not None with pytest.raises(RuntimeError, match=r"mstar\[metrics\]"): - metrics_module.ServingMetrics("dummy", in_flight=int) + metrics_module.ServingMetrics("dummy", in_flight=int, data_worker_queue_depth=int) with monkeypatch.context() as m: m.setitem(sys.modules, "prometheus_client.core", None) with pytest.raises(ModuleNotFoundError) as e: diff --git a/test/rust/test_rust_frontend.py b/test/rust/test_rust_frontend.py index 8ebbd33d10..cb1b58e9b9 100644 --- a/test/rust/test_rust_frontend.py +++ b/test/rust/test_rust_frontend.py @@ -253,7 +253,9 @@ def test_metrics_relay_the_api_server_registry_open_like_health(): from mstar.profile.format import RequestTiming with _running_stack(api_key="s3cret") as (port, stub, _bridge, _proc): - stub.metrics = ServingMetrics("qwen3_omni", in_flight=lambda: 3) + stub.metrics = ServingMetrics( + "qwen3_omni", in_flight=lambda: 3, data_worker_queue_depth=lambda: 0, + ) now = time.perf_counter() timing = RequestTiming(recv_time=now - 0.4, first_chunk_time=now - 0.3) stub.metrics.record(types.SimpleNamespace( @@ -926,7 +928,7 @@ def test_refusals_reach_the_api_server_registry(): key = {"Authorization": "Bearer s3cret"} gate = threading.Event() stub = _GatedKeyedStub(gate) - stub.metrics = ServingMetrics("qwen3_omni", in_flight=lambda: 0) + stub.metrics = ServingMetrics("qwen3_omni", in_flight=lambda: 0, data_worker_queue_depth=lambda: 0) env = {"MSTAR_API_KEY": "s3cret", "MSTAR_MAX_CONCURRENT_REQUESTS": "2", "MSTAR_MAX_BODY_MB": "1"} chat = json.dumps({"model": "qwen3_omni", "messages": [{"role": "user", "content": "hi"}]}).encode()