wehub-resource-sync
163 lines
5.2 KiB
Python
163 lines
5.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Tests for :class:`CallDispatcher` and the default v1 registrations."""
|
|
|
|
# Future
|
|
from __future__ import annotations
|
|
|
|
# Standard
|
|
from typing import Any
|
|
|
|
# Third Party
|
|
import pytest
|
|
|
|
# First Party
|
|
from lmcache.cli.commands.trace._dispatch import (
|
|
CallDispatcher,
|
|
ReplayContext,
|
|
build_default_dispatcher,
|
|
)
|
|
from lmcache.v1.distributed.api import ObjectKey
|
|
|
|
_SM_PREFIX = "lmcache.v1.distributed.storage_manager.StorageManager"
|
|
|
|
|
|
class _FakeSM:
|
|
"""Minimal StorageManager stand-in for dispatcher tests.
|
|
|
|
Records each call into ``self.calls`` so assertions can match
|
|
forwarded arguments exactly. For the context-manager entry, it
|
|
returns a tiny object whose ``__enter__``/``__exit__`` push into
|
|
``self.events`` so the FIFO ordering can be verified.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, dict[str, Any]]] = []
|
|
self.events: list[str] = []
|
|
|
|
def reserve_write(self, **kw: Any) -> dict[Any, Any]:
|
|
self.calls.append(("reserve_write", kw))
|
|
return {}
|
|
|
|
def finish_write(self, **kw: Any) -> None:
|
|
self.calls.append(("finish_write", kw))
|
|
|
|
def submit_prefetch_task(self, **kw: Any) -> None:
|
|
self.calls.append(("submit_prefetch_task", kw))
|
|
|
|
def finish_read_prefetched(self, **kw: Any) -> None:
|
|
self.calls.append(("finish_read_prefetched", kw))
|
|
|
|
def read_prefetched_results(self, keys: list[ObjectKey]) -> "_FakeCM":
|
|
return _FakeCM(self, keys)
|
|
|
|
|
|
class _FakeCM:
|
|
def __init__(self, parent: _FakeSM, keys: list[ObjectKey]) -> None:
|
|
self._parent = parent
|
|
self._keys = keys
|
|
|
|
def __enter__(self) -> None:
|
|
self._parent.events.append(f"enter-{len(self._keys)}")
|
|
|
|
def __exit__(self, *_exc: object) -> None:
|
|
self._parent.events.append(f"exit-{len(self._keys)}")
|
|
|
|
|
|
def _key(i: int) -> ObjectKey:
|
|
return ObjectKey(chunk_hash=bytes([i]), model_name="test", kv_rank=0)
|
|
|
|
|
|
class TestCallDispatcher:
|
|
def test_register_and_has(self):
|
|
d = CallDispatcher()
|
|
assert not d.has("x")
|
|
d.register("x", lambda c, a: None)
|
|
assert d.has("x")
|
|
|
|
def test_duplicate_registration_raises(self):
|
|
d = CallDispatcher()
|
|
d.register("x", lambda c, a: None)
|
|
with pytest.raises(ValueError):
|
|
d.register("x", lambda c, a: None)
|
|
|
|
def test_dispatch_unknown_qualname_raises_keyerror(self):
|
|
d = CallDispatcher()
|
|
with pytest.raises(KeyError):
|
|
d.dispatch("no.such.qualname", ReplayContext(sm=_FakeSM()), {})
|
|
|
|
def test_registered_qualnames_returns_new_list(self):
|
|
d = CallDispatcher()
|
|
d.register("a", lambda c, a: None)
|
|
qns = d.registered_qualnames()
|
|
qns.append("b")
|
|
assert "b" not in d.registered_qualnames()
|
|
|
|
|
|
class TestDefaultDispatcher:
|
|
def test_registers_all_v1_qualnames(self):
|
|
d = build_default_dispatcher()
|
|
expected = {
|
|
f"{_SM_PREFIX}.reserve_write",
|
|
f"{_SM_PREFIX}.finish_write",
|
|
f"{_SM_PREFIX}.submit_prefetch_task",
|
|
f"{_SM_PREFIX}.finish_read_prefetched",
|
|
f"{_SM_PREFIX}.read_prefetched_results.__enter__",
|
|
f"{_SM_PREFIX}.read_prefetched_results.__exit__",
|
|
}
|
|
assert set(d.registered_qualnames()) == expected
|
|
|
|
def test_simple_method_forwarded_with_kwargs(self):
|
|
sm = _FakeSM()
|
|
ctx = ReplayContext(sm=sm)
|
|
d = build_default_dispatcher()
|
|
d.dispatch(
|
|
f"{_SM_PREFIX}.reserve_write",
|
|
ctx,
|
|
{"keys": [_key(1)], "layout_desc": "LAYOUT", "mode": "new"},
|
|
)
|
|
assert sm.calls == [
|
|
(
|
|
"reserve_write",
|
|
{"keys": [_key(1)], "layout_desc": "LAYOUT", "mode": "new"},
|
|
),
|
|
]
|
|
|
|
def test_read_prefetched_enter_exit_fifo(self):
|
|
"""Two overlapping contexts with identical keys exit in FIFO order."""
|
|
sm = _FakeSM()
|
|
ctx = ReplayContext(sm=sm)
|
|
d = build_default_dispatcher()
|
|
|
|
keys = [_key(1), _key(2)]
|
|
enter = f"{_SM_PREFIX}.read_prefetched_results.__enter__"
|
|
exit_ = f"{_SM_PREFIX}.read_prefetched_results.__exit__"
|
|
|
|
d.dispatch(enter, ctx, {"keys": keys})
|
|
d.dispatch(enter, ctx, {"keys": keys})
|
|
assert sm.events == ["enter-2", "enter-2"]
|
|
# Both contexts share the same key tuple.
|
|
assert tuple(keys) in ctx.open_read_contexts
|
|
assert len(ctx.open_read_contexts[tuple(keys)]) == 2
|
|
|
|
d.dispatch(exit_, ctx, {"keys": keys})
|
|
assert sm.events == ["enter-2", "enter-2", "exit-2"]
|
|
d.dispatch(exit_, ctx, {"keys": keys})
|
|
assert sm.events == ["enter-2", "enter-2", "exit-2", "exit-2"]
|
|
# Fully drained → dict cleaned up.
|
|
assert ctx.open_read_contexts == {}
|
|
|
|
def test_exit_without_matching_enter_is_warning_only(self, caplog):
|
|
sm = _FakeSM()
|
|
ctx = ReplayContext(sm=sm)
|
|
d = build_default_dispatcher()
|
|
|
|
with caplog.at_level("WARNING"):
|
|
d.dispatch(
|
|
f"{_SM_PREFIX}.read_prefetched_results.__exit__",
|
|
ctx,
|
|
{"keys": [_key(1)]},
|
|
)
|
|
# No crash, no effect on sm.
|
|
assert sm.events == []
|