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()