wehub-resource-sync
448 lines
16 KiB
Python
448 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for full sync functionality (Tracker and KVController)."""
|
|
|
|
# Standard
|
|
import asyncio
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
|
|
# First Party
|
|
from lmcache.v1.cache_controller.commands.full_sync import FullSyncCommand
|
|
from lmcache.v1.cache_controller.controllers.registration_controller import (
|
|
RegistrationController,
|
|
)
|
|
from lmcache.v1.cache_controller.message import (
|
|
FullSyncStatusMsg,
|
|
HeartbeatMsg,
|
|
OpType,
|
|
)
|
|
from lmcache.v1.cache_controller.utils import FullSyncState, RegistryTree
|
|
|
|
# Local
|
|
# Test utilities from conftest
|
|
from .conftest import LOCATION, H
|
|
|
|
# ============= FullSyncTracker Tests =============
|
|
|
|
|
|
class TestFullSyncTracker:
|
|
"""Tests for FullSyncTracker."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"state,reason",
|
|
[
|
|
("none", "controller_restart"),
|
|
("syncing", None),
|
|
("completed", None),
|
|
("failed", "sync_failed_retry"),
|
|
],
|
|
)
|
|
def test_should_request_full_sync(self, state, reason):
|
|
"""Test sync request logic."""
|
|
tracker, reg = H.create_tracker()
|
|
H.reg_worker(reg, "i1", 0)
|
|
|
|
if state == "syncing":
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
elif state == "completed":
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
tracker.complete_sync("i1", 0, "s1", 100)
|
|
elif state == "failed":
|
|
tracker.set_need_full_sync_all(False)
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
tracker.mark_failed("i1", 0, "timeout")
|
|
|
|
need, r = tracker.should_request_full_sync("i1", 0)
|
|
assert (need, r) == (reason is not None, reason)
|
|
|
|
@pytest.mark.parametrize(
|
|
"action,syncing,state",
|
|
[
|
|
("none", False, None),
|
|
("start", True, FullSyncState.SYNCING),
|
|
("complete", False, FullSyncState.COMPLETED),
|
|
("fail", False, FullSyncState.FAILED),
|
|
],
|
|
)
|
|
def test_is_worker_syncing(self, action, syncing, state):
|
|
"""Test worker syncing status."""
|
|
tracker, reg = H.create_tracker()
|
|
H.reg_worker(reg, "i1", 0)
|
|
|
|
if action == "start":
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
elif action == "complete":
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
tracker.complete_sync("i1", 0, "s1", 100)
|
|
elif action == "fail":
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
tracker.mark_failed("i1", 0, "reason")
|
|
|
|
assert tracker.is_worker_syncing("i1", 0) == syncing
|
|
if action != "none":
|
|
assert reg.get_worker("i1", 0).sync_info.state == state
|
|
|
|
def test_start_sync_conflict(self):
|
|
"""Test sync start conflict handling."""
|
|
tracker, reg = H.create_tracker()
|
|
H.reg_worker(reg, "i1", 0)
|
|
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
assert tracker.start_sync("i1", 0, "s2", 200, 10) is False # Different ID
|
|
assert tracker.start_sync("i1", 0, "s1", 100, 5) is True # Same ID (retry)
|
|
|
|
@pytest.mark.parametrize("batches,keys_per", [(5, 20), (2, 50), (1, 100)])
|
|
def test_receive_batches(self, batches, keys_per):
|
|
"""Test batch receiving."""
|
|
tracker, reg = H.create_tracker()
|
|
H.reg_worker(reg, "i1", 0)
|
|
tracker.start_sync("i1", 0, "s1", 100, batches)
|
|
|
|
for i in range(batches):
|
|
assert tracker.receive_batch("i1", 0, "s1", i, keys_per) is True
|
|
|
|
info = reg.get_worker("i1", 0).sync_info
|
|
assert info.received_batches == set(range(batches))
|
|
assert info.received_keys_count == batches * keys_per
|
|
|
|
def test_receive_batch_wrong_sync_id(self):
|
|
"""Test batch with wrong sync ID rejected."""
|
|
tracker, reg = H.create_tracker()
|
|
H.reg_worker(reg, "i1", 0)
|
|
tracker.start_sync("i1", 0, "correct", 100, 5)
|
|
|
|
assert tracker.receive_batch("i1", 0, "wrong", 0, 20) is False
|
|
assert reg.get_worker("i1", 0).sync_info.received_keys_count == 0
|
|
|
|
def test_complete_sync(self):
|
|
"""Test sync completion."""
|
|
tracker, reg = H.create_tracker()
|
|
H.reg_worker(reg, "i1", 0)
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
|
|
for i in range(5):
|
|
tracker.receive_batch("i1", 0, "s1", i, 20)
|
|
|
|
assert tracker.complete_sync("i1", 0, "s1", 100) is True
|
|
assert reg.get_worker("i1", 0).sync_info.state == FullSyncState.COMPLETED
|
|
assert tracker.complete_sync("i1", 0, "wrong", 100) is False
|
|
|
|
def test_check_sync_timeout(self):
|
|
"""Test timeout detection."""
|
|
tracker, reg = H.create_tracker(timeout=0.1)
|
|
tracker.set_need_full_sync_all(False)
|
|
H.reg_worker(reg, "i1", 0)
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
|
|
time.sleep(0.15)
|
|
tracker.check_sync_timeout()
|
|
|
|
assert reg.get_worker("i1", 0).sync_info.state == FullSyncState.FAILED
|
|
need, reason = tracker.should_request_full_sync("i1", 0)
|
|
assert need and reason == "sync_failed_retry"
|
|
|
|
@pytest.mark.parametrize(
|
|
"total,completed,expected",
|
|
[(4, 0, 0.0), (4, 1, 0.25), (4, 2, 0.5), (4, 4, 1.0)],
|
|
)
|
|
def test_get_global_progress(self, total, completed, expected):
|
|
"""Test progress calculation."""
|
|
tracker, reg = H.create_tracker()
|
|
tracker.set_need_full_sync_all(False)
|
|
|
|
for i in range(total):
|
|
H.reg_worker(reg, "i1", i)
|
|
tracker.start_sync("i1", i, f"s{i}", 100, 5)
|
|
for i in range(completed):
|
|
tracker.complete_sync("i1", i, f"s{i}", 100)
|
|
|
|
assert abs(tracker.get_global_progress() - expected) < 0.001
|
|
|
|
@pytest.mark.parametrize(
|
|
"total,completed,threshold,can_exit",
|
|
[(4, 2, 0.5, True), (4, 1, 0.5, False), (4, 3, 0.8, False), (4, 4, 0.8, True)],
|
|
)
|
|
def test_can_exit_freeze(self, total, completed, threshold, can_exit):
|
|
"""Test freeze exit check."""
|
|
tracker, reg = H.create_tracker(threshold)
|
|
tracker.set_need_full_sync_all(True)
|
|
|
|
for i in range(total):
|
|
H.reg_worker(reg, "i1", i)
|
|
tracker.start_sync("i1", i, f"s{i}", 100, 5)
|
|
for i in range(completed):
|
|
tracker.complete_sync("i1", i, f"s{i}", 100)
|
|
|
|
assert tracker.can_exit_freeze() == can_exit
|
|
|
|
@pytest.mark.parametrize(
|
|
"received,missing",
|
|
[([], [0, 1, 2, 3, 4]), ([0, 2, 4], [1, 3]), ([0, 1, 2, 3, 4], [])],
|
|
)
|
|
def test_get_missing_batches(self, received, missing):
|
|
"""Test missing batches detection."""
|
|
tracker, reg = H.create_tracker()
|
|
H.reg_worker(reg, "i1", 0)
|
|
tracker.start_sync("i1", 0, "s1", 100, 5)
|
|
|
|
for b in received:
|
|
tracker.receive_batch("i1", 0, "s1", b, 20)
|
|
|
|
assert tracker.get_missing_batches("i1", 0, "s1") == missing
|
|
|
|
@pytest.mark.parametrize(
|
|
"setup,expected",
|
|
[
|
|
("no_sync", []), # No sync started
|
|
("wrong_id", []), # Wrong sync ID
|
|
("completed", []), # After completion
|
|
],
|
|
)
|
|
def test_get_missing_batches_edge_cases(self, setup, expected):
|
|
"""Test missing batches edge cases."""
|
|
tracker, reg = H.create_tracker()
|
|
H.reg_worker(reg, "i1", 0)
|
|
|
|
if setup == "wrong_id":
|
|
tracker.start_sync("i1", 0, "correct", 100, 5)
|
|
assert tracker.get_missing_batches("i1", 0, "wrong") == expected
|
|
elif setup == "completed":
|
|
tracker.start_sync("i1", 0, "correct", 100, 5)
|
|
tracker.complete_sync("i1", 0, "correct", 100)
|
|
assert tracker.get_missing_batches("i1", 0, "correct") == expected
|
|
else:
|
|
assert tracker.get_missing_batches("i1", 0, "s1") == expected
|
|
|
|
|
|
# ============= KVController Tests =============
|
|
|
|
|
|
class TestKVController:
|
|
"""Tests for KVController full sync handling."""
|
|
|
|
@pytest.fixture
|
|
def ctrl(self):
|
|
"""Create KVController."""
|
|
return H.create_controller()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_sync_start_clears_keys(self, ctrl):
|
|
"""Test sync start clears existing keys."""
|
|
H.reg_worker(ctrl.registry, "i1", 0)
|
|
for seq, key in enumerate([1, 2, 3]):
|
|
ctrl.registry.handle_batched_kv_operations(
|
|
H.batched_op("i1", 0, OpType.ADMIT, key, seq)
|
|
)
|
|
|
|
ret = await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s1", 100, 5))
|
|
|
|
assert ret.accepted is True
|
|
assert ctrl.full_sync_tracker.is_worker_syncing("i1", 0)
|
|
assert len(ctrl.registry.get_worker_kv_keys("i1", 0, LOCATION)) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_start_conflict(self, ctrl):
|
|
"""Test sync start conflict."""
|
|
H.reg_worker(ctrl.registry, "i1", 0)
|
|
|
|
r1 = await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s1", 100, 5))
|
|
r2 = await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s2", 200, 10))
|
|
r3 = await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s1", 100, 5))
|
|
|
|
assert r1.accepted and not r2.accepted and r3.accepted
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_handling(self, ctrl):
|
|
"""Test batch adds keys."""
|
|
H.reg_worker(ctrl.registry, "i1", 0)
|
|
await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s1", 15, 3))
|
|
|
|
for i in range(3):
|
|
await ctrl.handle_full_sync_batch(
|
|
H.batch_msg("i1", 0, "s1", i, list(range(i * 5, (i + 1) * 5)))
|
|
)
|
|
|
|
keys = ctrl.registry.get_worker_kv_keys("i1", 0, LOCATION)
|
|
assert len(keys) == 15 and all(i in keys for i in range(15))
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_wrong_sync_id_rejected(self, ctrl):
|
|
"""Test batch with wrong sync ID rejected."""
|
|
H.reg_worker(ctrl.registry, "i1", 0)
|
|
await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "correct", 10, 1))
|
|
|
|
keys_before = ctrl.registry.get_worker_kv_keys("i1", 0, LOCATION).copy()
|
|
await ctrl.handle_full_sync_batch(H.batch_msg("i1", 0, "wrong", 0, [1, 2, 3]))
|
|
|
|
assert ctrl.registry.get_worker_kv_keys("i1", 0, LOCATION) == keys_before
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_end_completes(self, ctrl):
|
|
"""Test sync end marks completion."""
|
|
H.reg_worker(ctrl.registry, "i1", 0)
|
|
await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s1", 5, 1))
|
|
await ctrl.handle_full_sync_batch(
|
|
H.batch_msg("i1", 0, "s1", 0, [1, 2, 3, 4, 5])
|
|
)
|
|
await ctrl.handle_full_sync_end(H.end_msg("i1", 0, "s1", 5))
|
|
|
|
assert not ctrl.full_sync_tracker.is_worker_syncing("i1", 0)
|
|
assert (
|
|
ctrl.registry.get_worker("i1", 0).sync_info.state == FullSyncState.COMPLETED
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_status(self, ctrl):
|
|
"""Test status query."""
|
|
ctrl.full_sync_tracker.completion_threshold = 0.5
|
|
for i in range(4):
|
|
H.reg_worker(ctrl.registry, "i1", i)
|
|
ctrl.full_sync_tracker.start_sync("i1", i, f"s{i}", 100, 5)
|
|
for i in range(2):
|
|
ctrl.full_sync_tracker.complete_sync("i1", i, f"s{i}", 100)
|
|
|
|
ret = await ctrl.handle_full_sync_status(FullSyncStatusMsg("i1", 0, "s0"))
|
|
|
|
assert (
|
|
ret.is_complete
|
|
and abs(ret.global_progress - 0.5) < 0.001
|
|
and ret.can_exit_freeze
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_status_missing_batches(self, ctrl):
|
|
"""Test status returns missing batches."""
|
|
H.reg_worker(ctrl.registry, "i1", 0)
|
|
await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s1", 100, 5))
|
|
await ctrl.handle_full_sync_batch(H.batch_msg("i1", 0, "s1", 0, [1, 2, 3]))
|
|
await ctrl.handle_full_sync_batch(H.batch_msg("i1", 0, "s1", 2, [7, 8, 9]))
|
|
|
|
ret = await ctrl.handle_full_sync_status(FullSyncStatusMsg("i1", 0, "s1"))
|
|
|
|
assert not ret.is_complete and ret.missing_batches == [1, 3, 4]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_incremental_discarded_during_sync(self, ctrl):
|
|
"""Test incremental messages discarded during sync."""
|
|
H.reg_worker(ctrl.registry, "i1", 0)
|
|
await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s1", 5, 1))
|
|
|
|
await ctrl.handle_batched_kv_operations(
|
|
H.batched_op("i1", 0, OpType.ADMIT, 888)
|
|
)
|
|
assert 888 not in ctrl.registry.get_worker_kv_keys("i1", 0, LOCATION)
|
|
|
|
await ctrl.handle_full_sync_batch(H.batch_msg("i1", 0, "s1", 0, [999]))
|
|
await ctrl.handle_batched_kv_operations(
|
|
H.batched_op("i1", 0, OpType.EVICT, 999)
|
|
)
|
|
assert 999 in ctrl.registry.get_worker_kv_keys("i1", 0, LOCATION)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_incremental_allowed_after_sync(self, ctrl):
|
|
"""Test incremental works after sync completion."""
|
|
H.reg_worker(ctrl.registry, "i1", 0)
|
|
await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s1", 0, 1))
|
|
await ctrl.handle_full_sync_end(H.end_msg("i1", 0, "s1", 0))
|
|
|
|
await ctrl.handle_batched_kv_operations(
|
|
H.batched_op("i1", 0, OpType.ADMIT, 1000, 1)
|
|
)
|
|
|
|
assert 1000 in ctrl.registry.get_worker_kv_keys("i1", 0, LOCATION)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_timeout(self):
|
|
"""Test sync timeout marks failed."""
|
|
ctrl = H.create_controller(timeout=0.1)
|
|
ctrl.full_sync_tracker.set_need_full_sync_all(False)
|
|
H.reg_worker(ctrl.registry, "i1", 0)
|
|
|
|
await ctrl.handle_full_sync_start(H.start_msg("i1", 0, "s1", 100, 5))
|
|
await asyncio.sleep(0.15)
|
|
ctrl.full_sync_tracker.check_sync_timeout()
|
|
|
|
assert ctrl.registry.get_worker("i1", 0).sync_info.state == FullSyncState.FAILED
|
|
need, reason = ctrl.full_sync_tracker.should_request_full_sync("i1", 0)
|
|
assert need and reason == "sync_failed_retry"
|
|
|
|
|
|
# ============= Integration Tests =============
|
|
|
|
|
|
class TestIntegration:
|
|
"""Integration tests for full sync flow."""
|
|
|
|
@pytest.fixture
|
|
def setup(self):
|
|
"""Setup controllers with shared registry."""
|
|
# Standard
|
|
from unittest.mock import MagicMock
|
|
|
|
# First Party
|
|
from lmcache.v1.cache_controller.controllers.kv_controller import KVController
|
|
|
|
reg = RegistryTree()
|
|
kv = KVController(reg, 0.5, 300.0)
|
|
kv.cluster_executor = MagicMock()
|
|
rc = RegistrationController()
|
|
rc.registry = reg
|
|
rc.kv_controller = kv
|
|
rc.cluster_executor = MagicMock()
|
|
return kv, rc
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_complete_sync_flow(self, setup):
|
|
"""Test complete sync flow for single worker."""
|
|
kv, rc = setup
|
|
|
|
await H.reg_worker_async(rc, "i1", 0)
|
|
for seq, key in enumerate([100, 200, 300]):
|
|
kv.registry.handle_batched_kv_operations(
|
|
H.batched_op("i1", 0, OpType.ADMIT, key, seq)
|
|
)
|
|
|
|
hb = await rc.heartbeat(HeartbeatMsg("i1", 0, "127.0.0.1", 8000, None))
|
|
assert len(hb.commands) == 1 and isinstance(hb.commands[0], FullSyncCommand)
|
|
|
|
await kv.handle_full_sync_start(H.start_msg("i1", 0, "s1", 10, 2))
|
|
assert len(kv.registry.get_worker_kv_keys("i1", 0, LOCATION)) == 0
|
|
|
|
await kv.handle_full_sync_batch(H.batch_msg("i1", 0, "s1", 0, [1, 2, 3, 4, 5]))
|
|
await kv.handle_full_sync_batch(H.batch_msg("i1", 0, "s1", 1, [6, 7, 8, 9, 10]))
|
|
await kv.handle_full_sync_end(H.end_msg("i1", 0, "s1", 10))
|
|
|
|
keys = kv.registry.get_worker_kv_keys("i1", 0, LOCATION)
|
|
assert keys == {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
|
|
|
|
hb2 = await rc.heartbeat(HeartbeatMsg("i1", 0, "127.0.0.1", 8000, None))
|
|
assert len(hb2.commands) == 0
|
|
|
|
status = await kv.handle_full_sync_status(FullSyncStatusMsg("i1", 0, "s1"))
|
|
assert status.is_complete and status.global_progress == 1.0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_progress_multiple_workers(self, setup):
|
|
"""Test sync progress with multiple workers."""
|
|
kv, rc = setup
|
|
workers = [(f"i{i}", i, f"192.168.1.{i + 1}") for i in range(4)]
|
|
|
|
for i in range(len(workers)):
|
|
inst, wid, ip = workers[i]
|
|
sid = f"s{i}"
|
|
await H.reg_worker_async(rc, inst, wid, ip)
|
|
await kv.handle_full_sync_start(H.start_msg(inst, wid, sid, 5, 1))
|
|
|
|
for i in range(2):
|
|
inst, wid, _ = workers[i]
|
|
sid = f"s{i}"
|
|
await kv.handle_full_sync_batch(
|
|
H.batch_msg(inst, wid, sid, 0, list(range(5)))
|
|
)
|
|
await kv.handle_full_sync_end(H.end_msg(inst, wid, sid, 5))
|
|
|
|
assert kv.full_sync_tracker.get_global_progress() == 0.5
|
|
assert kv.full_sync_tracker.can_exit_freeze() is True
|