wehub-resource-sync
540 lines
17 KiB
Python
540 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Tests for EventBus, EventSubscriber, and singleton management."""
|
|
|
|
# Standard
|
|
from unittest.mock import MagicMock, patch
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
|
|
# First Party
|
|
from lmcache.v1.mp_observability.event import Event, EventType
|
|
from lmcache.v1.mp_observability.event_bus import (
|
|
EventBus,
|
|
EventBusConfig,
|
|
EventSubscriber,
|
|
get_event_bus,
|
|
init_event_bus,
|
|
)
|
|
import lmcache.v1.mp_observability.event_bus as _bus_module
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_event(
|
|
event_type: EventType = EventType.L1_READ_FINISHED,
|
|
session_id: str = "s1",
|
|
**metadata,
|
|
) -> Event:
|
|
return Event(event_type=event_type, session_id=session_id, metadata=metadata)
|
|
|
|
|
|
class _RecordingSubscriber(EventSubscriber):
|
|
"""Test subscriber that records events."""
|
|
|
|
def __init__(self, event_types: list[EventType] | None = None):
|
|
if event_types is None:
|
|
event_types = [EventType.L1_READ_FINISHED]
|
|
self._event_types = event_types
|
|
self.events: list[Event] = []
|
|
self.shutdown_called = False
|
|
|
|
def get_subscriptions(self):
|
|
return {et: self._on_event for et in self._event_types}
|
|
|
|
def _on_event(self, event: Event) -> None:
|
|
self.events.append(event)
|
|
|
|
def shutdown(self) -> None:
|
|
self.shutdown_called = True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def restore_global_bus():
|
|
"""Save and restore the global singleton so tests don't leak state."""
|
|
saved = _bus_module._global_bus
|
|
yield
|
|
_bus_module._global_bus = saved
|
|
|
|
|
|
@pytest.fixture
|
|
def bus():
|
|
"""Enabled EventBus with small queue for fast tests."""
|
|
return EventBus(EventBusConfig(enabled=True, max_queue_size=100))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Subscribe / publish basics
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSubscription:
|
|
def test_subscribe_adds_callback(self, bus):
|
|
cb = MagicMock()
|
|
bus.subscribe(EventType.L1_READ_FINISHED, cb)
|
|
assert cb in bus._subscribers[EventType.L1_READ_FINISHED]
|
|
|
|
def test_multiple_callbacks_per_event_type(self, bus):
|
|
cb1, cb2 = MagicMock(), MagicMock()
|
|
bus.subscribe(EventType.L1_READ_FINISHED, cb1)
|
|
bus.subscribe(EventType.L1_READ_FINISHED, cb2)
|
|
assert len(bus._subscribers[EventType.L1_READ_FINISHED]) == 2
|
|
|
|
def test_register_subscriber(self, bus):
|
|
sub = _RecordingSubscriber()
|
|
bus.register_subscriber(sub)
|
|
assert sub in bus._registered_subscribers
|
|
assert len(bus._subscribers[EventType.L1_READ_FINISHED]) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Lifecycle
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLifecycle:
|
|
def test_start_launches_thread(self, bus):
|
|
bus.start()
|
|
assert bus._thread is not None
|
|
assert bus._thread.is_alive()
|
|
bus.stop()
|
|
|
|
def test_disabled_noop(self):
|
|
b = EventBus(EventBusConfig(enabled=False))
|
|
b.start()
|
|
assert b._thread is None
|
|
|
|
def test_thread_is_daemon(self, bus):
|
|
bus.start()
|
|
assert bus._thread.daemon is True
|
|
bus.stop()
|
|
|
|
def test_thread_name(self, bus):
|
|
bus.start()
|
|
assert bus._thread.name == "EventBus"
|
|
bus.stop()
|
|
|
|
def test_double_start_is_idempotent(self, bus):
|
|
bus.start()
|
|
thread = bus._thread
|
|
bus.start()
|
|
assert bus._thread is thread
|
|
bus.stop()
|
|
|
|
def test_double_stop(self, bus):
|
|
bus.start()
|
|
bus.stop()
|
|
bus.stop()
|
|
|
|
def test_stop_without_start(self, bus):
|
|
bus.stop()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Event dispatch
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestEventDispatch:
|
|
def test_event_reaches_subscriber(self, bus):
|
|
sub = _RecordingSubscriber()
|
|
bus.register_subscriber(sub)
|
|
bus.start()
|
|
|
|
bus.publish(_make_event())
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert len(sub.events) == 1
|
|
assert sub.events[0].event_type == EventType.L1_READ_FINISHED
|
|
|
|
def test_publish_stamps_timestamp(self, bus):
|
|
sub = _RecordingSubscriber()
|
|
bus.register_subscriber(sub)
|
|
bus.start()
|
|
|
|
before = time.time()
|
|
bus.publish(_make_event())
|
|
after = time.time()
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert len(sub.events) == 1
|
|
assert before <= sub.events[0].timestamp <= after
|
|
|
|
def test_events_dispatched_in_order(self, bus):
|
|
sub = _RecordingSubscriber()
|
|
bus.register_subscriber(sub)
|
|
bus.start()
|
|
|
|
for i in range(5):
|
|
bus.publish(_make_event(session_id=str(i)))
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert len(sub.events) == 5
|
|
for i, ev in enumerate(sub.events):
|
|
assert ev.session_id == str(i)
|
|
|
|
def test_event_only_dispatched_to_matching_subscribers(self, bus):
|
|
l1_sub = _RecordingSubscriber([EventType.L1_READ_FINISHED])
|
|
sm_sub = _RecordingSubscriber([EventType.SM_WRITE_RESERVED])
|
|
bus.register_subscriber(l1_sub)
|
|
bus.register_subscriber(sm_sub)
|
|
bus.start()
|
|
|
|
bus.publish(_make_event(event_type=EventType.L1_READ_FINISHED))
|
|
bus.publish(_make_event(event_type=EventType.SM_WRITE_RESERVED))
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert len(l1_sub.events) == 1
|
|
assert l1_sub.events[0].event_type == EventType.L1_READ_FINISHED
|
|
assert len(sm_sub.events) == 1
|
|
assert sm_sub.events[0].event_type == EventType.SM_WRITE_RESERVED
|
|
|
|
def test_multiple_subscribers_same_event(self, bus):
|
|
sub1 = _RecordingSubscriber()
|
|
sub2 = _RecordingSubscriber()
|
|
bus.register_subscriber(sub1)
|
|
bus.register_subscriber(sub2)
|
|
bus.start()
|
|
|
|
bus.publish(_make_event())
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert len(sub1.events) == 1
|
|
assert len(sub2.events) == 1
|
|
|
|
def test_metadata_preserved(self, bus):
|
|
sub = _RecordingSubscriber()
|
|
bus.register_subscriber(sub)
|
|
bus.start()
|
|
|
|
bus.publish(_make_event(count=42, name="test"))
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert sub.events[0].metadata["count"] == 42
|
|
assert sub.events[0].metadata["name"] == "test"
|
|
|
|
def test_disabled_bus_drops_events(self):
|
|
b = EventBus(EventBusConfig(enabled=False))
|
|
sub = _RecordingSubscriber()
|
|
b.register_subscriber(sub)
|
|
b.publish(_make_event())
|
|
assert len(b._queue) == 0
|
|
|
|
def test_late_registered_subscriber(self, bus):
|
|
bus.start()
|
|
time.sleep(0.05)
|
|
|
|
sub = _RecordingSubscriber()
|
|
bus.register_subscriber(sub)
|
|
|
|
bus.publish(_make_event())
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert len(sub.events) >= 1
|
|
|
|
def test_stop_drains_remaining_events(self, bus):
|
|
sub = _RecordingSubscriber()
|
|
bus.register_subscriber(sub)
|
|
# Don't start the drain thread — events accumulate in the queue
|
|
for _ in range(3):
|
|
bus.publish(_make_event())
|
|
assert len(bus._queue) == 3
|
|
|
|
# stop() should do a final drain
|
|
bus.stop()
|
|
assert len(sub.events) == 3
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Exception isolation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExceptionIsolation:
|
|
def test_bad_callback_doesnt_block_good(self, bus):
|
|
bad_cb = MagicMock(side_effect=RuntimeError("boom"))
|
|
bus.subscribe(EventType.L1_READ_FINISHED, bad_cb)
|
|
|
|
good_sub = _RecordingSubscriber()
|
|
bus.register_subscriber(good_sub)
|
|
bus.start()
|
|
|
|
bus.publish(_make_event())
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert len(good_sub.events) == 1
|
|
|
|
def test_shutdown_exception_isolated(self, bus):
|
|
bad_sub = _RecordingSubscriber()
|
|
bad_sub.shutdown = MagicMock(side_effect=RuntimeError("shutdown boom"))
|
|
good_sub = _RecordingSubscriber()
|
|
|
|
bus.register_subscriber(bad_sub)
|
|
bus.register_subscriber(good_sub)
|
|
bus.start()
|
|
bus.stop()
|
|
|
|
assert good_sub.shutdown_called
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Backpressure
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBackpressure:
|
|
def test_events_discarded_when_queue_full(self):
|
|
b = EventBus(EventBusConfig(enabled=True, max_queue_size=5))
|
|
# Don't start — nothing drains
|
|
for _ in range(10):
|
|
b.publish(_make_event())
|
|
|
|
assert len(b._queue) == 5
|
|
assert b._discard_count == 5
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Self-monitoring accessors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSelfMonitoring:
|
|
def test_queue_depth_reports_current_size(self, bus):
|
|
# Don't start the drain thread so events accumulate.
|
|
assert bus.queue_depth() == 0
|
|
for _ in range(3):
|
|
bus.publish(_make_event())
|
|
assert bus.queue_depth() == 3
|
|
|
|
def test_oldest_event_lag_zero_when_empty(self, bus):
|
|
assert bus.oldest_event_lag_seconds() == 0.0
|
|
|
|
def test_oldest_event_lag_grows_with_age(self, bus):
|
|
# Don't start drain — let events sit in the queue.
|
|
bus.publish(_make_event())
|
|
time.sleep(0.05)
|
|
lag = bus.oldest_event_lag_seconds()
|
|
assert lag >= 0.05
|
|
# Drain via stop() so the bus shuts down cleanly.
|
|
bus.stop()
|
|
|
|
def test_oldest_event_lag_returns_zero_after_drain(self, bus):
|
|
sub = _RecordingSubscriber()
|
|
bus.register_subscriber(sub)
|
|
bus.start()
|
|
bus.publish(_make_event())
|
|
time.sleep(0.15)
|
|
# Drained, queue should be empty.
|
|
assert bus.oldest_event_lag_seconds() == 0.0
|
|
bus.stop()
|
|
|
|
def test_dropped_events_count_starts_zero(self, bus):
|
|
assert bus.dropped_events_count() == 0
|
|
|
|
def test_dropped_events_count_increments_on_drop(self):
|
|
b = EventBus(EventBusConfig(enabled=True, max_queue_size=2))
|
|
# Don't start — nothing drains.
|
|
for _ in range(5):
|
|
b.publish(_make_event())
|
|
assert b.dropped_events_count() == 3
|
|
|
|
def test_subscriber_exception_counts_starts_empty(self, bus):
|
|
assert bus.subscriber_exception_counts() == {}
|
|
|
|
def test_subscriber_exception_counts_tracks_bound_methods(self, bus):
|
|
class _BadSub(EventSubscriber):
|
|
def get_subscriptions(self):
|
|
return {EventType.L1_READ_FINISHED: self._on_event}
|
|
|
|
def _on_event(self, event):
|
|
raise RuntimeError("boom")
|
|
|
|
bus.register_subscriber(_BadSub())
|
|
bus.start()
|
|
bus.publish(_make_event())
|
|
bus.publish(_make_event())
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
counts = bus.subscriber_exception_counts()
|
|
assert counts.get("_BadSub") == 2
|
|
|
|
def test_subscriber_exception_counts_returns_copy(self, bus):
|
|
class _BadSub(EventSubscriber):
|
|
def get_subscriptions(self):
|
|
return {EventType.L1_READ_FINISHED: self._on_event}
|
|
|
|
def _on_event(self, event):
|
|
raise RuntimeError("boom")
|
|
|
|
bus.register_subscriber(_BadSub())
|
|
bus.start()
|
|
bus.publish(_make_event())
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
snapshot = bus.subscriber_exception_counts()
|
|
snapshot["_BadSub"] = 99
|
|
assert bus.subscriber_exception_counts().get("_BadSub") == 1
|
|
|
|
def test_subscriber_exception_counts_uses_qualname_for_free_function(self, bus):
|
|
def free_callback(event):
|
|
raise RuntimeError("boom")
|
|
|
|
bus.subscribe(EventType.L1_READ_FINISHED, free_callback)
|
|
bus.start()
|
|
bus.publish(_make_event())
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
counts = bus.subscriber_exception_counts()
|
|
# Free functions are labeled by __qualname__, which here lives
|
|
# inside the test method.
|
|
assert any("free_callback" in name for name in counts)
|
|
|
|
def test_good_subscriber_not_counted(self, bus):
|
|
class _BadSub(EventSubscriber):
|
|
def get_subscriptions(self):
|
|
return {EventType.L1_READ_FINISHED: self._on_event}
|
|
|
|
def _on_event(self, event):
|
|
raise RuntimeError("boom")
|
|
|
|
good = _RecordingSubscriber()
|
|
bus.register_subscriber(_BadSub())
|
|
bus.register_subscriber(good)
|
|
bus.start()
|
|
bus.publish(_make_event())
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
counts = bus.subscriber_exception_counts()
|
|
assert "_BadSub" in counts
|
|
assert "_RecordingSubscriber" not in counts
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Global singleton
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGlobalSingleton:
|
|
def test_get_returns_instance(self):
|
|
bus = get_event_bus()
|
|
assert isinstance(bus, EventBus)
|
|
|
|
def test_init_replaces_global(self):
|
|
old = get_event_bus()
|
|
new = init_event_bus(EventBusConfig(enabled=False))
|
|
assert get_event_bus() is new
|
|
assert get_event_bus() is not old
|
|
|
|
def test_default_singleton_is_disabled(self):
|
|
bus = get_event_bus()
|
|
assert bus._config.enabled is False
|
|
|
|
def test_init_with_none_uses_defaults(self):
|
|
bus = init_event_bus()
|
|
assert bus._config.enabled is True
|
|
assert bus._config.max_queue_size == 10_000
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Block allocation event
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBlockAllocationEvent:
|
|
def test_publish_block_allocation_event(self, bus):
|
|
"""Verify MP_VLLM_BLOCK_ALLOCATION events are delivered to subscribers."""
|
|
sub = _RecordingSubscriber(event_types=[EventType.MP_VLLM_BLOCK_ALLOCATION])
|
|
bus.register_subscriber(sub)
|
|
bus.start()
|
|
|
|
# First Party
|
|
from lmcache.v1.multiprocess.custom_types import BlockAllocationRecord
|
|
|
|
records = [
|
|
BlockAllocationRecord(
|
|
req_id="req-1",
|
|
new_block_ids=[0, 1, 2],
|
|
new_token_ids=[10, 20, 30],
|
|
),
|
|
]
|
|
bus.publish(
|
|
_make_event(
|
|
event_type=EventType.MP_VLLM_BLOCK_ALLOCATION,
|
|
session_id="",
|
|
records=records,
|
|
)
|
|
)
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert len(sub.events) == 1
|
|
evt = sub.events[0]
|
|
assert evt.event_type == EventType.MP_VLLM_BLOCK_ALLOCATION
|
|
assert len(evt.metadata["records"]) == 1
|
|
assert evt.metadata["records"][0].req_id == "req-1"
|
|
assert evt.metadata["records"][0].new_block_ids == [0, 1, 2]
|
|
|
|
def test_block_allocation_not_delivered_to_other_subscriber(self, bus):
|
|
"""Verify block allocation events are not delivered to unrelated subscribers."""
|
|
sub = _RecordingSubscriber(event_types=[EventType.L1_READ_FINISHED])
|
|
bus.register_subscriber(sub)
|
|
bus.start()
|
|
|
|
bus.publish(
|
|
_make_event(
|
|
event_type=EventType.MP_VLLM_BLOCK_ALLOCATION,
|
|
session_id="",
|
|
)
|
|
)
|
|
time.sleep(0.15)
|
|
bus.stop()
|
|
|
|
assert len(sub.events) == 0
|
|
|
|
|
|
_REGISTER_GAUGE = "lmcache.v1.mp_observability.event_bus.register_gauge"
|
|
|
|
|
|
class TestSelfMonitoringGauges:
|
|
def test_init_registers_both_gauges(self):
|
|
with patch(_REGISTER_GAUGE) as mock_register:
|
|
EventBus(EventBusConfig(enabled=True))
|
|
names = [c.args[1] for c in mock_register.call_args_list]
|
|
assert set(names) == {
|
|
"lmcache_mp.event_bus.queue_depth",
|
|
"lmcache_mp.event_bus.drain_lag_seconds",
|
|
}
|
|
|
|
def test_queue_depth_callback_reflects_bus(self):
|
|
captured: dict[str, object] = {}
|
|
|
|
def _capture(_meter, name, _desc, func):
|
|
captured[name] = func
|
|
|
|
with patch(_REGISTER_GAUGE, side_effect=_capture):
|
|
bus = EventBus(EventBusConfig(enabled=True))
|
|
bus.publish(_make_event(session_id="s1"))
|
|
bus.publish(_make_event(session_id="s2"))
|
|
assert captured["lmcache_mp.event_bus.queue_depth"]() == 2
|