[Metrics] Add audio SLOs + cross-stage transfer families + per-(stage, replica) wrap for upstream vllm:* (#3576)

Signed-off-by: vraiti <vraiti@redhat.com>
Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
Co-authored-by: vraiti <vraiti@redhat.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -312,7 +312,7 @@ def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_inpu
     monkeypatch.setattr(
         engine_mod,
         "build_llm_stage_output_processor",
-        lambda plan, _cfg: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor,
+        lambda plan, _cfg, **_kw: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor,
     )
     monkeypatch.setattr(engine_mod, "build_stage0_input_processor", lambda _cfg: top_level_input_processor)
 

tests/engine/test_orchestrator.py

@@ -1157,3 +1157,37 @@ async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator
         assert stage0_r1.add_request_calls[1][0].request_id == "parent-neg"
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
+
+
+def test_orchestrator_does_not_re_introduce_global_stats_throttle() -> None:
+    """Regression: each (stage, replica) must independently publish its wrapped
+    vllm:* stats when its scheduler emits non-None scheduler_stats.
+
+    A previous version of Orchestrator carried a global self._last_stats_ts /
+    _stats_interval_s gate around _stat_logger.record(). Because
+    OmniSchedulerMixin.make_stats() already throttles at 1 Hz per scheduler
+    (one per (stage, replica)), the extra global gate starved every replica
+    other than the first to emit within each second — their {stage, replica}
+    gauges/counters went stale.
+
+    The fix removed the global gate entirely; the only signal needed is
+    'this replica's scheduler emitted non-None scheduler_stats'. This test
+    fails loudly if someone reintroduces the global throttle.
+    """
+    import inspect
+
+    from vllm_omni.engine.orchestrator import Orchestrator
+
+    source = inspect.getsource(Orchestrator)
+    assert "_last_stats_ts" not in source, (
+        "Orchestrator must not gate stat recording on a global timestamp. "
+        "OmniSchedulerMixin.make_stats() already throttles per scheduler "
+        "(per (stage, replica)); an outer global gate starves all but the "
+        "first replica to emit within each 1s window."
+    )
+    assert "_stats_interval_s" not in source
+    assert "raw_outputs.scheduler_stats is not None" in source, (
+        "Orchestrator must gate stat recording solely on "
+        "raw_outputs.scheduler_stats being non-None — the per-scheduler 1Hz "
+        "throttle in OmniSchedulerMixin.make_stats() is the only gate needed."
+    )

tests/metrics/test_modality.py

@@ -0,0 +1,399 @@
+from __future__ import annotations
+
+import pytest
+from prometheus_client import REGISTRY, generate_latest
+
+from vllm_omni.metrics import definitions as defs
+from vllm_omni.metrics.modality import (
+    OmniModalityMetrics,
+    observe_audio_first_packet,
+    observe_audio_streaming_finalize,
+    observe_modality_at_finalize,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+_MODEL = "test-modality-model"
+
+
+@pytest.fixture(scope="module")
+def mod() -> OmniModalityMetrics:
+    return OmniModalityMetrics(model_name=_MODEL)
+
+
+def _sample_value(output: str, line_prefix: str) -> float | None:
+    for line in output.splitlines():
+        if line.startswith(line_prefix):
+            return float(line.split()[-1])
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Family registration — 7 audio families served from modality.py
+# ---------------------------------------------------------------------------
+
+
+_EXPECTED_FAMILIES = [
+    defs.AUDIO_TTFP_S,
+    defs.AUDIO_DURATION_S,
+    defs.AUDIO_RTF_METRIC,
+    defs.AUDIO_FRAMES_METRIC,
+    defs.AUDIO_UNDERRUN_S,
+    defs.AUDIO_CONTINUITY_OK_METRIC,
+    defs.AUDIO_SKIPPED_REQUESTS_METRIC,
+]
+
+
+class TestRegistration:
+    def test_all_locked_families_present(self, mod: OmniModalityMetrics) -> None:
+        # Trigger at least one observation per family so the registry exposes them.
+        mod.observe_audio_ttfp("s", "r", 0.1)
+        mod.observe_audio_duration("s", "r", 1.0)
+        mod.observe_audio_rtf("s", "r", 0.5)
+        mod.inc_audio_frames("s", "r", 1)
+        mod.observe_audio_underrun("s", "r", 0.01)
+        mod.inc_audio_continuity_ok("s", "r", 100)
+        mod.inc_audio_skipped("s", "r", "malformed_codec")
+
+        out = generate_latest(REGISTRY).decode()
+        for name in _EXPECTED_FAMILIES:
+            assert f"# HELP {name}" in out, f"missing family: {name}"
+
+
+# ---------------------------------------------------------------------------
+# Audio observe API
+# ---------------------------------------------------------------------------
+
+
+class TestAudio:
+    def test_audio_ttfp_observed(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_ttfp", "0"
+        mod.observe_audio_ttfp(stage, replica, 0.42)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_TTFP_S}_count{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 1.0
+
+    def test_audio_duration_observed(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_dur", "0"
+        mod.observe_audio_duration(stage, replica, 3.5)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_DURATION_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 3.5
+
+    def test_audio_rtf_observed(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_rtf", "0"
+        mod.observe_audio_rtf(stage, replica, 0.45)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_RTF_METRIC}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 0.45
+
+    def test_audio_frames_inc(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_frames", "0"
+        mod.inc_audio_frames(stage, replica, 240)
+        mod.inc_audio_frames(stage, replica, 60)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_FRAMES_METRIC}_total{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 300.0
+
+    def test_audio_frames_zero_or_negative_skipped(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_zero", "0"
+        mod.inc_audio_frames(stage, replica, 0)
+        mod.inc_audio_frames(stage, replica, -5)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_FRAMES_METRIC}_total{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) is None
+
+    def test_audio_underrun_observed(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_under", "0"
+        mod.observe_audio_underrun(stage, replica, 0.05)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_UNDERRUN_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 0.05
+
+    def test_audio_underrun_negative_clamped_to_zero(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_under_neg", "0"
+        mod.observe_audio_underrun(stage, replica, -0.1)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_UNDERRUN_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 0.0
+
+    def test_audio_continuity_ok_carries_threshold_label(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_cont", "0"
+        mod.inc_audio_continuity_ok(stage, replica, threshold_ms=100)
+        mod.inc_audio_continuity_ok(stage, replica, threshold_ms=200)
+        out = generate_latest(REGISTRY).decode()
+        # Counter family auto-suffixes _total at exposition.
+        v100 = _sample_value(
+            out,
+            f'{defs.AUDIO_CONTINUITY_OK_METRIC}_total{{model_name="{_MODEL}",'
+            f'replica="{replica}",stage="{stage}",threshold_ms="100"}}',
+        )
+        v200 = _sample_value(
+            out,
+            f'{defs.AUDIO_CONTINUITY_OK_METRIC}_total{{model_name="{_MODEL}",'
+            f'replica="{replica}",stage="{stage}",threshold_ms="200"}}',
+        )
+        assert v100 == 1.0
+        assert v200 == 1.0
+
+    def test_audio_skipped_carries_reason_label(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_skip", "0"
+        mod.inc_audio_skipped(stage, replica, "malformed_codec")
+        mod.inc_audio_skipped(stage, replica, "")  # → "unknown"
+        out = generate_latest(REGISTRY).decode()
+        good = _sample_value(
+            out,
+            f'{defs.AUDIO_SKIPPED_REQUESTS_METRIC}_total{{model_name="{_MODEL}",'
+            f'reason="malformed_codec",replica="{replica}",stage="{stage}"}}',
+        )
+        unknown = _sample_value(
+            out,
+            f'{defs.AUDIO_SKIPPED_REQUESTS_METRIC}_total{{model_name="{_MODEL}",'
+            f'reason="unknown",replica="{replica}",stage="{stage}"}}',
+        )
+        assert good == 1.0
+        assert unknown == 1.0
+
+
+# ---------------------------------------------------------------------------
+# Stub-driven routing tests for the finalize / streaming helpers.
+# ---------------------------------------------------------------------------
+
+
+class _StubModMetrics:
+    def __init__(self):
+        self.calls: list[tuple] = []
+
+    def inc_audio_frames(self, s, r, n):
+        self.calls.append(("inc_audio_frames", s, r, n))
+
+    def observe_audio_duration(self, s, r, d):
+        self.calls.append(("observe_audio_duration", s, r, d))
+
+    def observe_audio_rtf(self, s, r, rtf):
+        self.calls.append(("observe_audio_rtf", s, r, rtf))
+
+    def observe_audio_underrun(self, s, r, u):
+        self.calls.append(("observe_audio_underrun", s, r, u))
+
+    def inc_audio_continuity_ok(self, s, r, threshold_ms):
+        self.calls.append(("inc_audio_continuity_ok", s, r, threshold_ms))
+
+    def inc_audio_skipped(self, s, r, reason):
+        self.calls.append(("inc_audio_skipped", s, r, reason))
+
+
+class _Bag:
+    def __init__(self, **kw):
+        self.__dict__.update(kw)
+
+
+class TestObserveModalityAtFinalize:
+    def test_audio_path_full(self):
+        stub = _StubModMetrics()
+        stage_metrics = _Bag(stage_gen_time_ms=500.0, audio_generated_frames=24000)
+        engine_outputs = _Bag(multimodal_output={"audio_sample_rate": 24000})
+
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=0,
+            stage_metrics=stage_metrics,
+            engine_outputs=engine_outputs,
+        )
+        # 24000 frames / 24000 Hz = 1.0s duration; gen 0.5s → rtf 0.5
+        assert ("inc_audio_frames", "1", "0", 24000) in stub.calls
+        assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
+        assert ("observe_audio_rtf", "1", "0", 0.5) in stub.calls
+        # Continuity/underrun NOT emitted from finalize — they come from the
+        # streaming hook because they need the per-chunk arrival timeline.
+        assert not any(c[0] == "observe_audio_underrun" for c in stub.calls)
+        assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls)
+
+    def test_audio_path_zero_frames_emits_skipped(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=0,
+            stage_metrics=_Bag(stage_gen_time_ms=300.0, audio_generated_frames=0),
+            engine_outputs=_Bag(multimodal_output={}),
+        )
+        assert ("inc_audio_frames", "1", "0", 0) in stub.calls
+        assert not any(c[0] == "observe_audio_duration" for c in stub.calls)
+        assert not any(c[0] == "observe_audio_rtf" for c in stub.calls)
+        assert ("inc_audio_skipped", "1", "0", "no_audio_data") in stub.calls
+
+    def test_audio_uses_resolved_sample_rate_from_multimodal_output(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=0,
+            stage_metrics=_Bag(stage_gen_time_ms=1000.0, audio_generated_frames=16000),
+            engine_outputs=_Bag(multimodal_output={"sample_rate": 16000}),
+        )
+        assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
+
+    def test_non_audio_output_type_skipped(self):
+        # Image / video / text output types are out of scope for this module.
+        stub = _StubModMetrics()
+        for output_type in ("text", "image", "video"):
+            observe_modality_at_finalize(
+                stub,
+                output_type=output_type,
+                stage_id=0,
+                replica_id=0,
+                stage_metrics=_Bag(stage_gen_time_ms=100.0),
+                engine_outputs=_Bag(),
+            )
+        assert stub.calls == []
+
+    def test_replica_id_none_skipped(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=None,
+            stage_metrics=_Bag(stage_gen_time_ms=500.0, audio_generated_frames=240),
+            engine_outputs=_Bag(multimodal_output={}),
+        )
+        assert stub.calls == []
+
+    def test_stage_metrics_none_skipped(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=0,
+            stage_metrics=None,
+            engine_outputs=_Bag(multimodal_output={}),
+        )
+        assert stub.calls == []
+
+
+class TestObserveAudioFirstPacket:
+    def test_observes_with_valid_inputs(self):
+        stub = _StubModMetrics()
+        stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t))
+
+        observe_audio_first_packet(
+            stub,
+            stage_id=1,
+            replica_id=0,
+            arrival_ts=100.0,
+            now_ts=100.42,
+        )
+        assert stub.calls == [("observe_audio_ttfp", "1", "0", pytest.approx(0.42))]
+
+    def test_replica_none_skipped(self):
+        stub = _StubModMetrics()
+        stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t))
+        observe_audio_first_packet(stub, stage_id=1, replica_id=None, arrival_ts=100.0, now_ts=100.5)
+        assert stub.calls == []
+
+    def test_arrival_ts_zero_skipped(self):
+        stub = _StubModMetrics()
+        stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t))
+        observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=0.0, now_ts=100.5)
+        assert stub.calls == []
+
+    def test_clock_skew_clamped_to_zero(self):
+        stub = _StubModMetrics()
+        stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t))
+        observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=100.5, now_ts=100.0)
+        assert stub.calls == [("observe_audio_ttfp", "1", "0", 0.0)]
+
+
+class TestObserveAudioStreamingFinalize:
+    def test_continuous_run_emits_underrun_and_continuity_ok(self):
+        stub = _StubModMetrics()
+        # Three back-to-back chunks of 0.5s each at 24 kHz s16le mono:
+        # 24000 Hz * 2 bytes * 0.5s = 24000 bytes per chunk.
+        arrivals = [0.5, 1.0, 1.5]
+        bytes_ = [24000, 24000, 24000]
+        observe_audio_streaming_finalize(
+            stub,
+            stage_id=1,
+            replica_id=0,
+            chunk_arrival_times_s=arrivals,
+            chunk_bytes=bytes_,
+            sample_rate=24000,
+            threshold_s=0.1,
+        )
+        # 0 underrun -> observe with 0.0, continuity_ok incremented at threshold 100.
+        assert ("observe_audio_underrun", "1", "0", 0.0) in stub.calls
+        assert ("inc_audio_continuity_ok", "1", "0", 100) in stub.calls
+
+    def test_late_chunk_emits_nonzero_underrun_and_no_continuity_inc(self):
+        stub = _StubModMetrics()
+        # Two chunks but second arrives 1s late; underrun > 0.1s threshold.
+        arrivals = [0.0, 2.0]
+        bytes_ = [24000, 24000]
+        observe_audio_streaming_finalize(
+            stub,
+            stage_id=1,
+            replica_id=0,
+            chunk_arrival_times_s=arrivals,
+            chunk_bytes=bytes_,
+            sample_rate=24000,
+            threshold_s=0.1,
+        )
+        underrun_calls = [c for c in stub.calls if c[0] == "observe_audio_underrun"]
+        assert underrun_calls
+        assert underrun_calls[0][-1] > 0.1
+        assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls)
+
+    def test_empty_arrivals_skipped(self):
+        stub = _StubModMetrics()
+        observe_audio_streaming_finalize(
+            stub,
+            stage_id=1,
+            replica_id=0,
+            chunk_arrival_times_s=[],
+            chunk_bytes=[],
+            sample_rate=24000,
+        )
+        assert stub.calls == []
+
+    def test_replica_none_skipped(self):
+        stub = _StubModMetrics()
+        observe_audio_streaming_finalize(
+            stub,
+            stage_id=1,
+            replica_id=None,
+            chunk_arrival_times_s=[0.5],
+            chunk_bytes=[24000],
+            sample_rate=24000,
+        )
+        assert stub.calls == []
+
+
+class TestBucketSelection:
+    def test_audio_rtf_uses_rtf_buckets(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_buckets", "0"
+        mod.observe_audio_rtf(stage, replica, 0.5)
+        out = generate_latest(REGISTRY).decode()
+        rtf_marker = f'{defs.AUDIO_RTF_METRIC}_bucket{{le="0.9"'
+        assert rtf_marker in out, "audio_rtf should use RTF_BUCKETS containing le=0.9"
+
+    def test_audio_ttfp_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_seconds", "0"
+        mod.observe_audio_ttfp(stage, replica, 0.1)
+        out = generate_latest(REGISTRY).decode()
+        sec_marker = f'{defs.AUDIO_TTFP_S}_bucket{{le="0.05"'
+        assert sec_marker in out, "audio_ttfp should use SECONDS_BUCKETS containing le=0.05"
+
+    def test_audio_underrun_uses_seconds_fast_buckets(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_fast", "0"
+        mod.observe_audio_underrun(stage, replica, 0.005)
+        out = generate_latest(REGISTRY).decode()
+        # SECONDS_FAST_BUCKETS includes le=0.001 which is absent from SECONDS_BUCKETS.
+        fast_marker = f'{defs.AUDIO_UNDERRUN_S}_bucket{{le="0.001"'
+        assert fast_marker in out, "audio_underrun_s should use SECONDS_FAST_BUCKETS containing le=0.001"

tests/metrics/test_prometheus.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+import re
+
+import pytest
+from prometheus_client import REGISTRY, CollectorRegistry, generate_latest
+
+from vllm_omni.metrics import OmniPrometheusMetrics
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODEL = "test-model"
+
+_PIPELINE_METRICS = [
+    "vllm:omni_num_requests_running",
+    "vllm:omni_num_requests_waiting",
+    "vllm:omni_requests_success",
+    "vllm:omni_e2e_request_latency_s",
+]
+
+
+@pytest.fixture(scope="module")
+def registry() -> CollectorRegistry:
+    return REGISTRY
+
+
+@pytest.fixture(scope="module")
+def prom() -> OmniPrometheusMetrics:
+    return OmniPrometheusMetrics(model_name=_MODEL)
+
+
+@pytest.fixture(scope="module")
+def scrape_output(prom: OmniPrometheusMetrics, registry: CollectorRegistry) -> str:
+    # Two natural completions (stop) + one length-cap + one failure (abort)
+    # exercise three distinct finished_reason buckets in the merged Counter.
+    prom.request_succeeded(e2e_seconds=1.5, finished_reason="stop")
+    prom.request_succeeded(e2e_seconds=2.0, finished_reason="stop")
+    prom.request_succeeded(e2e_seconds=3.0, finished_reason="length")
+    prom.request_failed()  # → finished_reason="abort"
+    prom.set_running(5)
+    prom.set_waiting(2)
+    return generate_latest(registry).decode()
+
+
+def _sample_value(output: str, metric_line: str) -> float | None:
+    for line in output.splitlines():
+        if line.startswith(metric_line):
+            return float(line.split()[-1])
+    return None
+
+
+class TestMetricObservation:
+    def test_all_metric_families_present(self, scrape_output: str) -> None:
+        for name in _PIPELINE_METRICS:
+            assert f"# HELP {name}" in scrape_output, f"missing metric family: {name}"
+
+    def test_counter_values(self, scrape_output: str) -> None:
+        # Per-reason buckets sourced from the merged completion Counter.
+        stop = _sample_value(
+            scrape_output,
+            f'vllm:omni_requests_success_total{{finished_reason="stop",model_name="{_MODEL}"}}',
+        )
+        assert stop == 2.0
+
+        length = _sample_value(
+            scrape_output,
+            f'vllm:omni_requests_success_total{{finished_reason="length",model_name="{_MODEL}"}}',
+        )
+        assert length == 1.0
+
+        abort = _sample_value(
+            scrape_output,
+            f'vllm:omni_requests_success_total{{finished_reason="abort",model_name="{_MODEL}"}}',
+        )
+        assert abort == 1.0
+
+    def test_gauge_values(self, scrape_output: str) -> None:
+        running = _sample_value(
+            scrape_output,
+            f'vllm:omni_num_requests_running{{model_name="{_MODEL}"}}',
+        )
+        assert running == 5.0
+
+        waiting = _sample_value(
+            scrape_output,
+            f'vllm:omni_num_requests_waiting{{model_name="{_MODEL}"}}',
+        )
+        assert waiting == 2.0
+
+    def test_histogram_counts(self, scrape_output: str) -> None:
+        # 3 successful completions (stop x2 + length x1) all observe e2e;
+        # the 1 failed completion only increments the Counter without
+        # observing the latency histogram, so the count stays at 3.
+        e2e_count = _sample_value(
+            scrape_output,
+            f'vllm:omni_e2e_request_latency_s_count{{model_name="{_MODEL}"}}',
+        )
+        assert e2e_count == 3.0
+
+
+class TestLabelCorrectness:
+    def test_pipeline_metrics_carry_model_name(self, scrape_output: str) -> None:
+        for name in _PIPELINE_METRICS:
+            pattern = rf'^{re.escape(name)}.*model_name="{re.escape(_MODEL)}"'
+            assert re.search(pattern, scrape_output, re.MULTILINE), f"{name} missing model_name label"
+
+    def test_no_legacy_engine_label(self, scrape_output: str) -> None:
+        assert 'engine="' not in scrape_output
+
+    def test_no_legacy_seconds_or_ms_families(self, scrape_output: str) -> None:
+        # Renamed: *_time_ms → *_s; *_time_seconds dropped.
+        for legacy in (
+            "vllm:omni_request_queue_time_seconds",
+            "vllm:omni_e2e_request_latency_seconds",
+        ):
+            assert legacy not in scrape_output, f"legacy family {legacy} still registered"
+
+
+class TestScrapeOutput:
+    def test_omni_metrics_in_default_registry(self, scrape_output: str) -> None:
+        for name in _PIPELINE_METRICS:
+            assert name in scrape_output
+
+    def test_process_metrics_in_default_registry(self, scrape_output: str) -> None:
+        # vllm:* metrics require a full PrometheusStatLogger with VllmConfig
+        # and are registered by the Orchestrator at server startup. Verifying
+        # their presence is covered by integration tests. Here we confirm the
+        # default registry is being scraped by checking for process_* metrics
+        # from the Python prometheus_client runtime.
+        assert "process_" in scrape_output
+
+
+class TestRequestLifecycleGauges:
+    """Regression tests for the running/waiting gauge race at request finalize.
+
+    Before the fix, the per-stage publish in OmniBase._process_single_result
+    ran while the completed request was still in self.request_states even
+    though the orchestrator had already decremented _running_counter — so a
+    single completed request briefly exposed running=0, total=1 → waiting=1,
+    which stuck on /metrics until another request triggered a refresh.
+
+    Two-layer fix: _process_single_result excludes the finalizing request
+    from `total`, and _log_summary_and_cleanup republishes gauges after the
+    pop as a fallback. These tests pin the post-cleanup state.
+    """
+
+    def test_running_and_waiting_zero_after_request_completes(self, registry: CollectorRegistry) -> None:
+        from types import SimpleNamespace
+
+        from vllm_omni.entrypoints.omni_base import OmniBase
+        from vllm_omni.metrics.prometheus import OmniPrometheusMetrics, OmniRequestCounter
+
+        obj = object.__new__(OmniBase)
+        obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter())
+        obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test")
+        obj.request_states = {}
+        obj.log_stats = False
+
+        # Simulate request lifecycle: start (counter 0→1, dict {} → {req}),
+        # then finalize (counter 1→0, dict still holds the request because
+        # _log_summary_and_cleanup hasn't run yet — this is the exact race
+        # window the fix addresses).
+        obj.engine._running_counter.increment()
+        obj.request_states["req-1"] = SimpleNamespace(metrics=SimpleNamespace(e2e_done={"req-1"}))
+        obj.engine._running_counter.decrement()
+
+        obj._log_summary_and_cleanup("req-1")
+
+        assert obj.engine._running_counter.value == 0
+        assert len(obj.request_states) == 0
+        out = generate_latest(registry).decode()
+        assert _sample_value(out, 'vllm:omni_num_requests_running{model_name="lifecycle-test"}') == 0.0
+        assert _sample_value(out, 'vllm:omni_num_requests_waiting{model_name="lifecycle-test"}') == 0.0
+
+    def test_gauges_reflect_remaining_requests_after_one_completes(self, registry: CollectorRegistry) -> None:
+        from types import SimpleNamespace
+
+        from vllm_omni.entrypoints.omni_base import OmniBase
+        from vllm_omni.metrics.prometheus import OmniPrometheusMetrics, OmniRequestCounter
+
+        obj = object.__new__(OmniBase)
+        obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter())
+        obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-2")
+        obj.request_states = {}
+        obj.log_stats = False
+
+        # Two in flight, one finalizes — running should report 1, waiting 0.
+        obj.engine._running_counter.increment()
+        obj.engine._running_counter.increment()
+        obj.request_states["req-1"] = SimpleNamespace(metrics=SimpleNamespace(e2e_done={"req-1"}))
+        obj.request_states["req-2"] = SimpleNamespace(metrics=SimpleNamespace(e2e_done=set()))
+        obj.engine._running_counter.decrement()
+
+        obj._log_summary_and_cleanup("req-1")
+
+        out = generate_latest(registry).decode()
+        assert _sample_value(out, 'vllm:omni_num_requests_running{model_name="lifecycle-test-2"}') == 1.0
+        assert _sample_value(out, 'vllm:omni_num_requests_waiting{model_name="lifecycle-test-2"}') == 0.0

tests/metrics/test_stat_logger.py

@@ -0,0 +1,456 @@
+from __future__ import annotations
+
+import pytest
+from prometheus_client import CollectorRegistry, generate_latest
+
+from vllm_omni.metrics.stat_logger import (
+    _ENGINE_INDEX_MAP,
+    OmniPrometheusStatLogger,
+    _engine_to_stage_replica,
+    _OmniKVConnectorProm,
+    _OmniPerfMetricsProm,
+    _OmniSpecDecodingProm,
+    _RelabelCounter,
+    _RelabelGauge,
+    _RelabelHistogram,
+    _rewrite_labelnames,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _isolate_engine_map():
+    """Each test gets a clean _ENGINE_INDEX_MAP."""
+    _ENGINE_INDEX_MAP.clear()
+    yield
+    _ENGINE_INDEX_MAP.clear()
+
+
+@pytest.fixture
+def registry() -> CollectorRegistry:
+    return CollectorRegistry()
+
+
+# ---------------------------------------------------------------------------
+# _rewrite_labelnames
+# ---------------------------------------------------------------------------
+
+
+class TestRewriteLabelnames:
+    def test_engine_at_end(self):
+        assert _rewrite_labelnames(["model_name", "engine"]) == [
+            "model_name",
+            "stage",
+            "replica",
+        ]
+
+    def test_engine_in_middle(self):
+        # Upstream uses `labelnames + ["reason"]` etc., putting engine in middle.
+        assert _rewrite_labelnames(["model_name", "engine", "reason"]) == [
+            "model_name",
+            "stage",
+            "replica",
+            "reason",
+        ]
+
+    def test_no_engine_label(self):
+        # Unaffected (e.g. omni's own families that don't use engine).
+        assert _rewrite_labelnames(["model_name"]) == ["model_name"]
+
+    def test_tuple_input_returns_tuple(self):
+        out = _rewrite_labelnames(("model_name", "engine"))
+        assert isinstance(out, tuple)
+        assert out == ("model_name", "stage", "replica")
+
+    def test_none_passthrough(self):
+        assert _rewrite_labelnames(None) is None
+
+
+# ---------------------------------------------------------------------------
+# _engine_to_stage_replica
+# ---------------------------------------------------------------------------
+
+
+class TestEngineToStageReplica:
+    def test_int_engine_value(self):
+        # Mirrors upstream `.labels(engine=idx, ...)` with int (loggers.py:510).
+        _ENGINE_INDEX_MAP[7] = ("talker", "1")
+        assert _engine_to_stage_replica(7) == ("talker", "1")
+
+    def test_str_engine_value(self):
+        # Mirrors upstream `metrics_info["engine"] = str(idx)` (loggers.py:1055).
+        _ENGINE_INDEX_MAP[2] = ("thinker", "0")
+        assert _engine_to_stage_replica("2") == ("thinker", "0")
+
+    def test_missing_engine_idx_raises(self):
+        # Empty map → fail-fast rather than emit a wrong (stage, replica).
+        with pytest.raises(KeyError):
+            _engine_to_stage_replica(999)
+
+
+# ---------------------------------------------------------------------------
+# Wrapper class behavior
+# ---------------------------------------------------------------------------
+
+
+class TestRelabelGauge:
+    def test_labelnames_rewritten_at_creation(self, registry):
+        g = _RelabelGauge(
+            name="omni_test_gauge",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        assert g._labelnames == ("model_name", "stage", "replica")
+
+    def test_labels_kwarg_translated(self, registry):
+        _ENGINE_INDEX_MAP[5] = ("diffusion", "0")
+        g = _RelabelGauge(
+            name="omni_test_gauge_kwarg",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        g.labels(engine=5, model_name="qwen-omni").set(42.0)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_gauge_kwarg{model_name="qwen-omni",replica="0",stage="diffusion"} 42.0' in out
+
+    def test_labels_positional_passthrough(self, registry):
+        # Double-rewrite guard: the per_engine_labelvalues setter rewrites
+        # the values to 3-tuple [model_name, stage, replica] BEFORE
+        # create_metric_per_engine fans them into .labels(*values). The mixin
+        # must detect that args length already matches the rewritten
+        # labelnames and pass through, otherwise it would re-interpret
+        # args[engine_label_index] as an engine_idx and splice (stage, replica)
+        # again, blowing label count to 4.
+        g = _RelabelGauge(
+            name="omni_test_gauge_pos",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        g.labels("qwen-omni", "thinker", "0").set(7.0)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_gauge_pos{model_name="qwen-omni",replica="0",stage="thinker"} 7.0' in out
+
+    def test_multiprocess_mode_kwarg_passthrough(self, registry):
+        # Upstream creates Gauges with multiprocess_mode="mostrecent" — must not
+        # be eaten by our mixin.
+        g = _RelabelGauge(
+            name="omni_test_gauge_mp",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            multiprocess_mode="mostrecent",
+            registry=registry,
+        )
+        assert g._multiprocess_mode == "mostrecent"
+
+
+class TestRelabelCounter:
+    def test_labelnames_rewritten(self, registry):
+        c = _RelabelCounter(
+            name="omni_test_counter",
+            documentation="test",
+            labelnames=["model_name", "engine", "finished_reason"],
+            registry=registry,
+        )
+        assert c._labelnames == (
+            "model_name",
+            "stage",
+            "replica",
+            "finished_reason",
+        )
+
+    def test_labels_kwarg_translated(self, registry):
+        _ENGINE_INDEX_MAP[2] = ("thinker", "0")
+        c = _RelabelCounter(
+            name="omni_test_counter_kwarg",
+            documentation="test",
+            labelnames=["model_name", "engine", "finished_reason"],
+            registry=registry,
+        )
+        c.labels(engine=2, model_name="m", finished_reason="stop").inc(3)
+
+        out = generate_latest(registry).decode()
+        assert (
+            'omni_test_counter_kwarg_total{finished_reason="stop",model_name="m",replica="0",stage="thinker"} 3.0'
+            in out
+        )
+
+
+class TestRelabelHistogram:
+    def test_labelnames_rewritten(self, registry):
+        h = _RelabelHistogram(
+            name="omni_test_histo",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        assert h._labelnames == ("model_name", "stage", "replica")
+
+    def test_labels_kwarg_translated_and_observe(self, registry):
+        _ENGINE_INDEX_MAP[0] = ("talker", "0")
+        h = _RelabelHistogram(
+            name="omni_test_histo_obs",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        h.labels(engine=0, model_name="m").observe(0.5)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_histo_obs_count{model_name="m",replica="0",stage="talker"} 1.0' in out
+
+    def test_no_engine_label_unaffected(self, registry):
+        # Families without engine label (e.g. omni-side own metrics) pass through.
+        h = _RelabelHistogram(
+            name="omni_test_no_engine",
+            documentation="test",
+            labelnames=["model_name"],
+            registry=registry,
+        )
+        assert h._labelnames == ("model_name",)
+        h.labels(model_name="m").observe(1.0)
+
+
+# ---------------------------------------------------------------------------
+# Positional .labels() with engine value (loggers.py:646, 679)
+# ---------------------------------------------------------------------------
+
+
+class TestPositionalEngine:
+    def test_positional_engine_at_middle_index(self, registry):
+        # Mirrors `counter_prompt_tokens_by_source.labels(model_name, str(idx), source)`.
+        # Family original labelnames = ["model_name", "engine", "source"].
+        _ENGINE_INDEX_MAP[5] = ("talker", "0")
+        c = _RelabelCounter(
+            name="omni_test_pos_mid",
+            documentation="test",
+            labelnames=["model_name", "engine", "source"],
+            registry=registry,
+        )
+        c.labels("m", "5", "decoder").inc(2)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_pos_mid_total{model_name="m",replica="0",source="decoder",stage="talker"} 2.0' in out
+
+    def test_positional_engine_with_int_value(self, registry):
+        # Defensive: positional form may also receive an int (we accept both).
+        _ENGINE_INDEX_MAP[3] = ("thinker", "1")
+        c = _RelabelCounter(
+            name="omni_test_pos_int",
+            documentation="test",
+            labelnames=["model_name", "engine", "reason"],
+            registry=registry,
+        )
+        c.labels("m", 3, "stop").inc()
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_pos_int_total{model_name="m",reason="stop",replica="1",stage="thinker"} 1.0' in out
+
+
+# ---------------------------------------------------------------------------
+# String-form engine kwarg (loggers.py:1056 info_gauge)
+# ---------------------------------------------------------------------------
+
+
+class TestStrEngineKwarg:
+    def test_engine_kwarg_str_form(self, registry):
+        # Mirrors `info_gauge.labels(**metrics_info)` where metrics_info["engine"]="0".
+        _ENGINE_INDEX_MAP[0] = ("thinker", "0")
+        g = _RelabelGauge(
+            name="omni_test_info",
+            documentation="test",
+            labelnames=["cache_size", "engine"],
+            multiprocess_mode="mostrecent",
+            registry=registry,
+        )
+        # Upstream pattern: pass everything as kwargs from the metrics_info dict.
+        g.labels(cache_size="big", engine="0").set(1)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_info{cache_size="big",replica="0",stage="thinker"} 1.0' in out
+
+
+# ---------------------------------------------------------------------------
+# Child metric does not re-trigger relabel logic
+# ---------------------------------------------------------------------------
+
+
+class TestChildNoRecursion:
+    def test_child_set_does_not_relookup(self, registry):
+        # Once .labels() returns a child, subsequent .set()/.inc() must not
+        # consult _ENGINE_INDEX_MAP again. We verify by clearing the map
+        # AFTER labels() and proving .set() still works.
+        _ENGINE_INDEX_MAP[4] = ("diffusion", "0")
+        g = _RelabelGauge(
+            name="omni_test_child",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        child = g.labels(engine=4, model_name="m")
+        _ENGINE_INDEX_MAP.clear()  # would break a second .labels() lookup
+        child.set(99.0)  # but set() is on the bound child — no map needed
+
+        # Re-populate so generate_latest doesn't trip on anything else.
+        _ENGINE_INDEX_MAP[4] = ("diffusion", "0")
+        out = generate_latest(registry).decode()
+        assert 'omni_test_child{model_name="m",replica="0",stage="diffusion"} 99.0' in out
+
+
+# ---------------------------------------------------------------------------
+# OmniPrometheusStatLogger — focused on the wrap mechanics (full
+# PrometheusStatLogger init requires a real VllmConfig and is exercised by
+# the orchestrator integration tests).
+# ---------------------------------------------------------------------------
+
+
+class TestOmniPrometheusStatLogger:
+    def test_class_slots_point_to_wrappers(self):
+        # Upstream's __init__ uses self._gauge_cls(...) etc. when constructing
+        # families; class-level slot override is how we inject the relabel logic.
+        assert OmniPrometheusStatLogger._gauge_cls is _RelabelGauge
+        assert OmniPrometheusStatLogger._counter_cls is _RelabelCounter
+        assert OmniPrometheusStatLogger._histogram_cls is _RelabelHistogram
+
+    def test_per_engine_labelvalues_setter_rewrites_to_3tuple(self):
+        # Construct via __new__ to skip the upstream PrometheusStatLogger __init__
+        # (which needs a real VllmConfig). We only verify the property descriptor.
+        sl = OmniPrometheusStatLogger.__new__(OmniPrometheusStatLogger)
+        sl._stage_replica_map = {
+            0: ("thinker", "0"),
+            1: ("talker", "0"),
+            2: ("talker", "1"),
+        }
+
+        # Mirror upstream's loggers.py:433 assignment shape.
+        sl.per_engine_labelvalues = {
+            0: ["my-model", "0"],
+            1: ["my-model", "1"],
+            2: ["my-model", "2"],
+        }
+
+        # Getter should return the 3-tuple form for downstream
+        # create_metric_per_engine consumers.
+        assert sl.per_engine_labelvalues == {
+            0: ["my-model", "thinker", "0"],
+            1: ["my-model", "talker", "0"],
+            2: ["my-model", "talker", "1"],
+        }
+
+    def test_per_engine_labelvalues_getter_returns_internal_dict(self):
+        sl = OmniPrometheusStatLogger.__new__(OmniPrometheusStatLogger)
+        sl._stage_replica_map = {0: ("thinker", "0")}
+        sl._omni_per_engine_labelvalues = {0: ["m", "thinker", "0"]}
+        assert sl.per_engine_labelvalues == {0: ["m", "thinker", "0"]}
+
+    def test_stage_replica_map_property_exposed(self):
+        sl = OmniPrometheusStatLogger.__new__(OmniPrometheusStatLogger)
+        srm = {0: ("thinker", "0"), 1: ("diffusion", "0")}
+        sl._stage_replica_map = srm
+        assert sl.stage_replica_map is srm
+
+    def test_init_populates_engine_index_map(self):
+        # Simulate the bookkeeping portion of __init__ (clear + update) without
+        # calling super, since super needs a real VllmConfig.
+        _ENGINE_INDEX_MAP[99] = ("stale", "stale")  # leftover from prior
+        srm = {0: ("thinker", "0"), 1: ("talker", "0")}
+
+        # Manually invoke the bookkeeping the way __init__ does it.
+        _ENGINE_INDEX_MAP.clear()
+        _ENGINE_INDEX_MAP.update(srm)
+
+        assert dict(_ENGINE_INDEX_MAP) == srm
+        assert 99 not in _ENGINE_INDEX_MAP  # old entry was cleared
+
+
+# ---------------------------------------------------------------------------
+# Helper-class wraps for upstream's spec_decoding / kv_connector /
+# perf_metrics sub-collectors. Without these, OmniPrometheusStatLogger
+# crashes at startup with `Incorrect label count` because each helper builds
+# its internal Counter/Gauge/Histogram families with raw 2-element labelnames
+# (passed via constructor arg) while consuming the rewritten 3-element
+# per_engine_labelvalues from the property descriptor.
+# ---------------------------------------------------------------------------
+
+
+class TestHelperClassWraps:
+    def test_perf_metrics_wrap_routes_through_relabel_counter(self):
+        assert _OmniPerfMetricsProm._counter_cls is _RelabelCounter
+
+    def test_spec_decoding_wrap_routes_through_relabel_counter(self):
+        assert _OmniSpecDecodingProm._counter_cls is _RelabelCounter
+
+    def test_kv_connector_wrap_routes_through_all_three_relabel_classes(self):
+        # KVConnector lets each connector build any of Gauge/Counter/Histogram,
+        # so all three slots must be intercepted.
+        assert _OmniKVConnectorProm._gauge_cls is _RelabelGauge
+        assert _OmniKVConnectorProm._counter_cls is _RelabelCounter
+        assert _OmniKVConnectorProm._histogram_cls is _RelabelHistogram
+
+    def test_omni_logger_slots_point_to_helper_subclasses(self):
+        # Upstream's PrometheusStatLogger.__init__ instantiates each sub-helper
+        # via `self._<name>_cls(...)`, so the slot overrides on the omni
+        # subclass are what routes through to the relabel mixin.
+        assert OmniPrometheusStatLogger._perf_metrics_cls is _OmniPerfMetricsProm
+        assert OmniPrometheusStatLogger._spec_decoding_cls is _OmniSpecDecodingProm
+        assert OmniPrometheusStatLogger._kv_connector_cls is _OmniKVConnectorProm
+
+
+# ---------------------------------------------------------------------------
+# Double-rewrite guard. The mixin's positional-args path used to
+# unconditionally splice (stage, replica) at engine_label_index. Because
+# per_engine_labelvalues is rewritten to 3-tuples *before* feeding into
+# create_metric_per_engine, that splice would run a second time on the
+# already-rewritten values and blow the label count to 4. The guard now
+# detects len(args) == len(self._labelnames) and short-circuits to passthrough.
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def fresh_registry():
+    return CollectorRegistry()
+
+
+class TestDoubleRewriteGuard:
+    def test_pre_rewritten_3tuple_passes_through(self, fresh_registry):
+        # 2-label original → 3-label rewritten family. Caller passes 3 values
+        # (the rewritten shape) and they should land verbatim, not get
+        # re-spliced.
+        _ENGINE_INDEX_MAP.clear()
+        _ENGINE_INDEX_MAP[0] = ("0", "0")
+        _ENGINE_INDEX_MAP[1] = ("1", "0")
+        g = _RelabelGauge(
+            name="dr_pre_rewritten",
+            documentation="t",
+            labelnames=["model_name", "engine"],
+            registry=fresh_registry,
+        )
+        # 3 positional args matching the rewritten 3-label family.
+        g.labels("m", "1", "0").set(42)
+        out = generate_latest(fresh_registry).decode()
+        assert 'dr_pre_rewritten{model_name="m",replica="0",stage="1"} 42.0' in out
+
+    def test_legacy_2tuple_with_extra_label_still_splices(self, fresh_registry):
+        # 3-label original (engine in middle) → 4-label rewritten family.
+        # Caller passes 3 values matching the ORIGINAL labelnames (the
+        # gauge_waiting_by_reason / counter_request_success pattern from
+        # upstream loggers.py:646, 679). The mixin must splice
+        # (stage, replica) at engine's position to reach the 4-label family.
+        _ENGINE_INDEX_MAP.clear()
+        _ENGINE_INDEX_MAP[1] = ("1", "0")
+        c = _RelabelCounter(
+            name="dr_legacy_with_extra",
+            documentation="t",
+            labelnames=["model_name", "engine", "reason"],
+            registry=fresh_registry,
+        )
+        # 3 positional args matching ORIGINAL labelnames (model_name,
+        # engine_str, reason).
+        c.labels("m", "1", "stop").inc(3)
+        out = generate_latest(fresh_registry).decode()
+        assert 'dr_legacy_with_extra_total{model_name="m",reason="stop",replica="0",stage="1"} 3.0' in out

tests/metrics/test_transfer.py

@@ -0,0 +1,346 @@
+from __future__ import annotations
+
+import pytest
+from prometheus_client import REGISTRY, generate_latest
+
+from vllm_omni.metrics import definitions as defs
+from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats, StageStats
+from vllm_omni.metrics.transfer import OmniTransferMetrics
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+_MODEL = "test-transfer-model"
+
+
+@pytest.fixture(scope="module")
+def tx() -> OmniTransferMetrics:
+    return OmniTransferMetrics(model_name=_MODEL)
+
+
+def _sample_value(output: str, line_prefix: str) -> float | None:
+    for line in output.splitlines():
+        if line.startswith(line_prefix):
+            return float(line.split()[-1])
+    return None
+
+
+_EXPECTED_FAMILIES = [
+    defs.TRANSFER_SIZE_BYTES,
+    defs.TRANSFER_TX_S,
+    defs.TRANSFER_RX_S,
+    defs.TRANSFER_IN_FLIGHT_S,
+]
+
+
+# ---------------------------------------------------------------------------
+# Family registration
+# ---------------------------------------------------------------------------
+
+
+class TestRegistration:
+    def test_all_four_families_present(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_size(0, 0, 1, 0, 1024)
+        tx.observe_tx_time(0, 0, 1, 0, 0.005)
+        tx.observe_rx_time(0, 0, 1, 0, 0.008)
+        tx.observe_in_flight_time(0, 0, 1, 0, 0.002)
+
+        out = generate_latest(REGISTRY).decode()
+        for name in _EXPECTED_FAMILIES:
+            assert f"# HELP {name}" in out, f"missing family: {name}"
+
+    def test_legacy_ms_family_names_dropped(self, tx: OmniTransferMetrics) -> None:
+        out = generate_latest(REGISTRY).decode()
+        for legacy in (
+            "vllm:omni_transfer_tx_time_ms",
+            "vllm:omni_transfer_rx_decode_time_ms",
+            "vllm:omni_transfer_in_flight_time_ms",
+        ):
+            assert legacy not in out, f"legacy family {legacy} still registered"
+
+
+# ---------------------------------------------------------------------------
+# Observe APIs
+# ---------------------------------------------------------------------------
+
+
+class TestObserveSize:
+    def test_size_observed_with_correct_labels(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_size(2, 0, 3, 1, 65536)
+        out = generate_latest(REGISTRY).decode()
+        prefix = (
+            f"{defs.TRANSFER_SIZE_BYTES}_sum"
+            f'{{from_replica="0",from_stage="2",model_name="{_MODEL}",'
+            f'to_replica="1",to_stage="3"}}'
+        )
+        assert _sample_value(out, prefix) == 65536.0
+
+
+class TestObserveTxTime:
+    def test_tx_time_observed_in_seconds(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_tx_time(2, 1, 3, 0, 0.0125)
+        out = generate_latest(REGISTRY).decode()
+        prefix = (
+            f"{defs.TRANSFER_TX_S}_sum"
+            f'{{from_replica="1",from_stage="2",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="3"}}'
+        )
+        assert _sample_value(out, prefix) == pytest.approx(0.0125)
+
+
+class TestObserveRxTime:
+    def test_rx_time_observed_in_seconds(self, tx: OmniTransferMetrics) -> None:
+        # Use a (from, to) combo not shared with TestRegistration's setup
+        # observe on the rx_s Histogram, otherwise the per-test sum picks
+        # up accumulated values from earlier observes.
+        tx.observe_rx_time(3, 0, 4, 0, 0.0042)
+        out = generate_latest(REGISTRY).decode()
+        prefix = (
+            f"{defs.TRANSFER_RX_S}_sum"
+            f'{{from_replica="0",from_stage="3",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="4"}}'
+        )
+        assert _sample_value(out, prefix) == pytest.approx(0.0042)
+
+
+class TestObserveInFlightTime:
+    def test_in_flight_time_observed_in_seconds(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_in_flight_time(3, 0, 4, 0, 0.0017)
+        out = generate_latest(REGISTRY).decode()
+        prefix = (
+            f"{defs.TRANSFER_IN_FLIGHT_S}_sum"
+            f'{{from_replica="0",from_stage="3",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="4"}}'
+        )
+        assert _sample_value(out, prefix) == pytest.approx(0.0017)
+
+
+# ---------------------------------------------------------------------------
+# Multi (from, to) cardinality
+# ---------------------------------------------------------------------------
+
+
+class TestCardinality:
+    def test_multiple_edges_produce_independent_series(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_size(5, 0, 6, 0, 100)
+        tx.observe_size(5, 0, 6, 1, 200)
+        tx.observe_size(5, 1, 6, 0, 300)
+
+        out = generate_latest(REGISTRY).decode()
+
+        prefix_a = (
+            f"{defs.TRANSFER_SIZE_BYTES}_sum"
+            f'{{from_replica="0",from_stage="5",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="6"}}'
+        )
+        prefix_b = (
+            f"{defs.TRANSFER_SIZE_BYTES}_sum"
+            f'{{from_replica="0",from_stage="5",model_name="{_MODEL}",'
+            f'to_replica="1",to_stage="6"}}'
+        )
+        prefix_c = (
+            f"{defs.TRANSFER_SIZE_BYTES}_sum"
+            f'{{from_replica="1",from_stage="5",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="6"}}'
+        )
+        assert _sample_value(out, prefix_a) == 100.0
+        assert _sample_value(out, prefix_b) == 200.0
+        assert _sample_value(out, prefix_c) == 300.0
+
+
+# ---------------------------------------------------------------------------
+# Bucket selection
+# ---------------------------------------------------------------------------
+
+
+class TestBucketSelection:
+    def test_size_uses_bytes_buckets(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_size(7, 0, 8, 0, 4096)
+        out = generate_latest(REGISTRY).decode()
+        # prometheus_client renders bucket boundary as float (e.g. ``1024.0``)
+        # regardless of the declared int in BYTES_BUCKETS.
+        marker = f'{defs.TRANSFER_SIZE_BYTES}_bucket{{from_replica="0",from_stage="7",le="1024.0"'
+        assert marker in out
+
+    def test_time_families_use_seconds_fast_buckets(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_tx_time(7, 0, 8, 0, 0.001)
+        out = generate_latest(REGISTRY).decode()
+        # SECONDS_FAST_BUCKETS has le=0.001 — distinctive vs SECONDS_BUCKETS (no 0.001).
+        marker = f'{defs.TRANSFER_TX_S}_bucket{{from_replica="0",from_stage="7",le="0.001"'
+        assert marker in out
+
+
+# ---------------------------------------------------------------------------
+# OrchestratorAggregator emit hook
+# ---------------------------------------------------------------------------
+
+
+class _StubTransferEmitter:
+    """Records every observe_* call so the hook routing can be asserted
+    without standing up a Prometheus registry."""
+
+    def __init__(self) -> None:
+        self.calls: list[tuple] = []
+
+    def observe_size(self, fs, fr, ts, tr, n):
+        self.calls.append(("observe_size", fs, fr, ts, tr, n))
+
+    def observe_tx_time(self, fs, fr, ts, tr, t):
+        self.calls.append(("observe_tx_time", fs, fr, ts, tr, t))
+
+    def observe_rx_time(self, fs, fr, ts, tr, t):
+        self.calls.append(("observe_rx_time", fs, fr, ts, tr, t))
+
+    def observe_in_flight_time(self, fs, fr, ts, tr, t):
+        self.calls.append(("observe_in_flight_time", fs, fr, ts, tr, t))
+
+
+def _make_stats(stage_id, request_id, *, rx_decode=0.0, rx_in_flight=0.0, rx_bytes=0):
+    """Minimal StageRequestStats for record_transfer_rx input."""
+    return StageRequestStats(
+        batch_id=1,
+        batch_size=1,
+        num_tokens_in=0,
+        num_tokens_out=0,
+        stage_gen_time_ms=0.0,
+        rx_transfer_bytes=rx_bytes,
+        rx_decode_time_ms=rx_decode,
+        rx_in_flight_time_ms=rx_in_flight,
+        stage_stats=StageStats(),
+        stage_id=stage_id,
+        request_id=request_id,
+    )
+
+
+class TestEmitHookTx:
+    def test_record_transfer_tx_emits_size_and_tx_time_in_seconds(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: {0: 1, 1: 0}.get(s),
+        )
+        agg.record_transfer_tx(
+            from_stage=0,
+            to_stage=1,
+            request_id="r-tx-1",
+            size_bytes=2048,
+            tx_time_ms=7.5,
+            used_shm=False,
+        )
+        # tx_time_ms=7.5 -> 0.0075s at the emit boundary.
+        assert emitter.calls == [
+            ("observe_size", 0, 1, 1, 0, 2048),
+            ("observe_tx_time", 0, 1, 1, 0, pytest.approx(0.0075)),
+        ]
+
+    def test_record_transfer_tx_no_emit_when_emitter_none(self):
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=None,
+            replica_resolver=lambda s, rid: 0,
+        )
+        evt = agg.record_transfer_tx(
+            from_stage=0,
+            to_stage=1,
+            request_id="r-tx-2",
+            size_bytes=128,
+            tx_time_ms=1.0,
+            used_shm=True,
+        )
+        # Underlying accumulator keeps ms; only the Prometheus emit converts.
+        assert evt is not None
+        assert evt.size_bytes == 128
+        assert evt.tx_time_ms == 1.0
+
+    def test_record_transfer_tx_no_emit_when_resolver_returns_none(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: None,
+        )
+        agg.record_transfer_tx(
+            from_stage=0,
+            to_stage=1,
+            request_id="r-tx-3",
+            size_bytes=512,
+            tx_time_ms=2.0,
+            used_shm=False,
+        )
+        assert emitter.calls == []
+
+    def test_record_transfer_tx_no_emit_when_one_side_resolves_to_none(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: 0 if s == 0 else None,
+        )
+        agg.record_transfer_tx(
+            from_stage=0,
+            to_stage=1,
+            request_id="r-tx-4",
+            size_bytes=64,
+            tx_time_ms=0.5,
+            used_shm=False,
+        )
+        assert emitter.calls == []
+
+
+class TestEmitHookRx:
+    def test_record_transfer_rx_emits_rx_and_in_flight_in_seconds(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: {0: 1, 1: 0}.get(s),
+        )
+        stats = _make_stats(stage_id=1, request_id="r-rx-1", rx_decode=4.2, rx_in_flight=1.7)
+        agg.record_transfer_rx(stats)
+        assert emitter.calls == [
+            ("observe_rx_time", 0, 1, 1, 0, pytest.approx(0.0042)),
+            ("observe_in_flight_time", 0, 1, 1, 0, pytest.approx(0.0017)),
+        ]
+
+    def test_record_transfer_rx_skips_stage_zero(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: 0,
+        )
+        stats = _make_stats(stage_id=0, request_id="r-rx-2", rx_decode=4.2)
+        agg.record_transfer_rx(stats)
+        assert emitter.calls == []
+
+    def test_record_transfer_rx_no_emit_when_emitter_none(self):
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+        )
+        stats = _make_stats(stage_id=1, request_id="r-rx-3", rx_decode=1.0, rx_in_flight=0.5)
+        evt = agg.record_transfer_rx(stats)
+        assert evt is not None
+        assert evt.rx_decode_time_ms == 1.0
+        assert evt.in_flight_time_ms == 0.5