134 lines
4.3 KiB
Python
134 lines
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import json
|
|
import time
|
|
|
|
import pytest
|
|
|
|
import vllm_omni.engine.orchestrator_monitor as monitor_mod
|
|
from vllm_omni.engine.orchestrator_monitor import (
|
|
_NULL_ORCH_MONITOR,
|
|
OrchestratorMonitor,
|
|
create_orch_monitor,
|
|
replica_key,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_replica_key_format():
|
|
assert replica_key(0, 0) == "stage=0,replica=0"
|
|
assert replica_key(2, 1) == "stage=2,replica=1"
|
|
|
|
|
|
def test_create_orch_monitor_disabled_returns_null_singleton():
|
|
monitor = create_orch_monitor(enabled=False, replica_sampler=lambda: {})
|
|
assert monitor is _NULL_ORCH_MONITOR
|
|
assert create_orch_monitor(enabled=False, replica_sampler=lambda: {}) is monitor
|
|
monitor.note_loop(idle=True)
|
|
monitor.flush()
|
|
|
|
|
|
def test_note_loop_and_flush_json(tmp_path, monkeypatch):
|
|
out_path = tmp_path / "orch_monitor.json"
|
|
monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))
|
|
|
|
now = [0.0]
|
|
|
|
def fake_monotonic() -> float:
|
|
return now[0]
|
|
|
|
monkeypatch.setattr(time, "monotonic", fake_monotonic)
|
|
|
|
def sampler() -> dict[str, tuple[int, int]]:
|
|
return {replica_key(0, 0): (5, 10)}
|
|
|
|
monitor = create_orch_monitor(enabled=True, replica_sampler=sampler)
|
|
monitor.register_replica(0, 0)
|
|
|
|
monitor.note_loop(idle=False)
|
|
monitor.set_dispatch_queue_size(9)
|
|
monitor.set_dispatch_queue_size(7)
|
|
monitor.note_loop(idle=True)
|
|
|
|
now[0] = 1.5
|
|
monitor.flush()
|
|
|
|
payload = json.loads(out_path.read_text())
|
|
assert payload["configured_window_s"] == monitor_mod._WINDOW_S
|
|
assert payload["windows"]["loop_idle"] == [1]
|
|
assert payload["windows"]["loop_active"] == [1]
|
|
assert payload["windows"]["dispatch_queue_size"] == [7]
|
|
assert payload["windows"]["dispatch_queue_high_water"] == [9]
|
|
assert payload["replicas"]["stage=0,replica=0"]["outputs_queue_size"] == [5]
|
|
assert payload["replicas"]["stage=0,replica=0"]["inflight"] == [10]
|
|
|
|
|
|
@pytest.mark.parametrize("drain", [False, True])
|
|
def test_dispatch_backlog_carries_across_windows(tmp_path, monkeypatch, drain):
|
|
out_path = tmp_path / "orch_monitor.json"
|
|
monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))
|
|
now = [0.0]
|
|
monkeypatch.setattr(time, "monotonic", lambda: now[0])
|
|
monitor = OrchestratorMonitor(replica_sampler=lambda: {})
|
|
monitor.set_dispatch_queue_size(3)
|
|
monitor.set_dispatch_queue_size(1)
|
|
|
|
now[0] = monitor_mod._WINDOW_S
|
|
monitor.note_loop(idle=False)
|
|
if drain:
|
|
monitor.set_dispatch_queue_size(0)
|
|
now[0] += monitor_mod._WINDOW_S
|
|
monitor.note_loop(idle=False)
|
|
now[0] += monitor_mod._WINDOW_S
|
|
monitor.flush()
|
|
|
|
windows = json.loads(out_path.read_text())["windows"]
|
|
remaining = 0 if drain else 1
|
|
assert windows["dispatch_queue_size"] == [1, remaining, remaining]
|
|
assert windows["dispatch_queue_high_water"] == [3, 1, remaining]
|
|
|
|
|
|
def test_register_replica_backfills_prior_windows(tmp_path, monkeypatch):
|
|
out_path = tmp_path / "orch_monitor.json"
|
|
monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))
|
|
|
|
now = [0.0]
|
|
monkeypatch.setattr(time, "monotonic", lambda: now[0])
|
|
|
|
monitor = OrchestratorMonitor(replica_sampler=lambda: {})
|
|
monitor.note_loop(idle=True)
|
|
now[0] = 1.5
|
|
monitor.register_replica(1, 0)
|
|
monitor.note_loop(idle=False)
|
|
monitor.flush()
|
|
|
|
series = json.loads(out_path.read_text())["replicas"]["stage=1,replica=0"]
|
|
assert series["outputs_queue_size"] == [0, 0]
|
|
assert series["inflight"] == [0, 0]
|
|
|
|
|
|
def test_note_loop_survives_sampler_failure():
|
|
def bad_sampler() -> dict[str, tuple[int, int]]:
|
|
raise RuntimeError("sampler failed")
|
|
|
|
monitor = OrchestratorMonitor(replica_sampler=bad_sampler)
|
|
monitor.register_replica(0, 0)
|
|
monitor.note_loop(idle=True)
|
|
monitor.flush()
|
|
|
|
|
|
def test_flush_is_idempotent(tmp_path, monkeypatch):
|
|
out_path = tmp_path / "orch_monitor.json"
|
|
monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))
|
|
|
|
monitor = OrchestratorMonitor(replica_sampler=lambda: {})
|
|
monitor.register_replica(0, 0)
|
|
monitor.note_loop(idle=False)
|
|
monitor.flush()
|
|
monitor.flush()
|
|
|
|
payload = json.loads(out_path.read_text())
|
|
assert len(payload["windows"]["duration_s"]) == 1
|