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