wehub-resource-sync
257 lines
9.8 KiB
Python
257 lines
9.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Unit tests for AsyncSerdeProcessor.
|
|
|
|
Verifies the async, notifier-based contract:
|
|
- submit returns a task id immediately.
|
|
- The corresponding notifier fd is signaled on completion.
|
|
- query_result returns the bool outcome, None before completion and after
|
|
being consumed (non-idempotent).
|
|
- Failing serialize/deserialize produces result=False and still signals fd.
|
|
- Serialize and deserialize use distinct event fds.
|
|
"""
|
|
|
|
# Standard
|
|
from typing import Callable, Optional
|
|
import select
|
|
import time
|
|
|
|
# First Party
|
|
from lmcache.v1.distributed.api import MemoryLayoutDesc
|
|
from lmcache.v1.distributed.serde import (
|
|
AsyncSerdeProcessor,
|
|
Deserializer,
|
|
Serializer,
|
|
)
|
|
from lmcache.v1.platform import consume_fd
|
|
|
|
|
|
class _FakeSerializer(Serializer):
|
|
def __init__(self, transform: Optional[Callable[[int], None]] = None) -> None:
|
|
self._transform = transform
|
|
self.calls = 0
|
|
|
|
def serialize(self, src, dst) -> int: # type: ignore[no-untyped-def]
|
|
if self._transform is not None:
|
|
self._transform(self.calls)
|
|
self.calls += 1
|
|
return 1
|
|
|
|
def estimate_serialized_size(self, layout_desc: MemoryLayoutDesc) -> int:
|
|
return 1
|
|
|
|
|
|
class _FakeDeserializer(Deserializer):
|
|
def __init__(self, transform: Optional[Callable[[int], None]] = None) -> None:
|
|
self._transform = transform
|
|
self.calls = 0
|
|
|
|
def deserialize(self, src, dst) -> None: # type: ignore[no-untyped-def]
|
|
if self._transform is not None:
|
|
self._transform(self.calls)
|
|
self.calls += 1
|
|
|
|
|
|
def _wait_for_fd(fd: int, timeout_s: float = 2.0) -> bool:
|
|
"""Wait until ``fd`` is readable or timeout. Drains the pending signal."""
|
|
poller = select.poll()
|
|
poller.register(fd, select.POLLIN)
|
|
deadline = time.monotonic() + timeout_s
|
|
while time.monotonic() < deadline:
|
|
remaining_ms = int(max(0, (deadline - time.monotonic()) * 1000))
|
|
if poller.poll(remaining_ms):
|
|
try:
|
|
consume_fd(fd)
|
|
except OSError:
|
|
pass
|
|
return True
|
|
return False
|
|
|
|
|
|
def test_serialize_and_deserialize_fds_are_distinct() -> None:
|
|
processor = AsyncSerdeProcessor(_FakeSerializer(), _FakeDeserializer())
|
|
try:
|
|
assert (
|
|
processor.get_serialize_event_fd() != processor.get_deserialize_event_fd()
|
|
)
|
|
finally:
|
|
processor.close()
|
|
|
|
|
|
def test_serialize_signals_fd_and_result_is_true() -> None:
|
|
serializer = _FakeSerializer()
|
|
processor = AsyncSerdeProcessor(serializer, _FakeDeserializer())
|
|
try:
|
|
task_id = processor.submit_serialize([object()], [object()]) # type: ignore[list-item]
|
|
assert _wait_for_fd(processor.get_serialize_event_fd()), "fd never signaled"
|
|
assert processor.query_serialize_result(task_id) is True
|
|
# Non-idempotent: second query returns None.
|
|
assert processor.query_serialize_result(task_id) is None
|
|
assert serializer.calls == 1
|
|
finally:
|
|
processor.close()
|
|
|
|
|
|
def test_deserialize_signals_fd_and_result_is_true() -> None:
|
|
deserializer = _FakeDeserializer()
|
|
processor = AsyncSerdeProcessor(_FakeSerializer(), deserializer)
|
|
try:
|
|
task_id = processor.submit_deserialize([object()], [object()]) # type: ignore[list-item]
|
|
assert _wait_for_fd(processor.get_deserialize_event_fd()), "fd never signaled"
|
|
assert processor.query_deserialize_result(task_id) is True
|
|
assert processor.query_deserialize_result(task_id) is None
|
|
assert deserializer.calls == 1
|
|
finally:
|
|
processor.close()
|
|
|
|
|
|
def test_serialize_failure_reports_false() -> None:
|
|
"""If the sync serializer raises, query_serialize_result returns False."""
|
|
|
|
def _boom(_i: int) -> None:
|
|
raise RuntimeError("serialize failed")
|
|
|
|
processor = AsyncSerdeProcessor(_FakeSerializer(_boom), _FakeDeserializer())
|
|
try:
|
|
task_id = processor.submit_serialize([object()], [object()]) # type: ignore[list-item]
|
|
assert _wait_for_fd(processor.get_serialize_event_fd()), "fd never signaled"
|
|
assert processor.query_serialize_result(task_id) is False
|
|
finally:
|
|
processor.close()
|
|
|
|
|
|
def test_query_returns_none_before_completion() -> None:
|
|
"""Querying an unknown/not-yet-completed task yields None, not an error."""
|
|
processor = AsyncSerdeProcessor(_FakeSerializer(), _FakeDeserializer())
|
|
try:
|
|
# No task submitted with id 42 — should be None, not raise.
|
|
assert processor.query_serialize_result(42) is None
|
|
assert processor.query_deserialize_result(42) is None
|
|
finally:
|
|
processor.close()
|
|
|
|
|
|
def test_estimate_serialized_size_delegates_to_serializer() -> None:
|
|
serializer = _FakeSerializer()
|
|
processor = AsyncSerdeProcessor(serializer, _FakeDeserializer())
|
|
try:
|
|
layout = MemoryLayoutDesc(shapes=[], dtypes=[])
|
|
assert processor.estimate_serialized_size(layout) == 1
|
|
finally:
|
|
processor.close()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# set_used_size propagation: after serialize returns ``n``, the processor
|
|
# narrows the destination's logical size so downstream L2 adapters that
|
|
# read ``obj.get_size()`` / ``obj.byte_array`` see exactly ``n`` bytes.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _RecordingSerializer(Serializer):
|
|
"""A serializer that always claims to have written ``n_bytes`` so
|
|
tests can assert the processor narrows the destination accordingly.
|
|
"""
|
|
|
|
def __init__(self, n_bytes: int) -> None:
|
|
self._n_bytes = n_bytes
|
|
|
|
def serialize(self, src, dst) -> int: # type: ignore[no-untyped-def]
|
|
return self._n_bytes
|
|
|
|
def estimate_serialized_size(self, layout_desc: MemoryLayoutDesc) -> int:
|
|
return self._n_bytes
|
|
|
|
|
|
class _NarrowableDst:
|
|
"""A minimal MemoryObj-like destination that records every
|
|
``set_used_size`` call. Not a real ``MemoryObj`` -- the test only
|
|
needs the part of the interface the processor touches."""
|
|
|
|
def __init__(self) -> None:
|
|
self.used_sizes: list[int] = []
|
|
|
|
def set_used_size(self, n: int) -> None:
|
|
self.used_sizes.append(n)
|
|
|
|
|
|
def test_serialize_narrows_dst_to_actual_bytes_written() -> None:
|
|
"""After a successful serialize that returns ``n``, the processor
|
|
calls ``dst.set_used_size(n)`` so the over-allocated destination
|
|
buffer reports the actual written length, not its allocated size."""
|
|
serializer = _RecordingSerializer(n_bytes=213)
|
|
processor = AsyncSerdeProcessor(serializer, _FakeDeserializer())
|
|
try:
|
|
dst = _NarrowableDst()
|
|
task_id = processor.submit_serialize([object()], [dst]) # type: ignore[list-item]
|
|
assert _wait_for_fd(processor.get_serialize_event_fd()), "fd never signaled"
|
|
assert processor.query_serialize_result(task_id) is True
|
|
assert dst.used_sizes == [213]
|
|
finally:
|
|
processor.close()
|
|
|
|
|
|
def test_serialize_failure_does_not_narrow_dst() -> None:
|
|
"""If the serializer raises, the processor must NOT narrow the
|
|
destination -- failure cleanup typically reuses the original
|
|
allocation and the next attempt depends on the layout-derived
|
|
size."""
|
|
|
|
class _BoomSerializer(Serializer):
|
|
def serialize(self, src, dst) -> int: # type: ignore[no-untyped-def]
|
|
raise RuntimeError("serialize failed")
|
|
|
|
def estimate_serialized_size(self, layout_desc: MemoryLayoutDesc) -> int:
|
|
return 1
|
|
|
|
processor = AsyncSerdeProcessor(_BoomSerializer(), _FakeDeserializer())
|
|
try:
|
|
dst = _NarrowableDst()
|
|
task_id = processor.submit_serialize([object()], [dst]) # type: ignore[list-item]
|
|
assert _wait_for_fd(processor.get_serialize_event_fd()), "fd never signaled"
|
|
assert processor.query_serialize_result(task_id) is False
|
|
assert dst.used_sizes == [], "narrowing must not happen on failure"
|
|
finally:
|
|
processor.close()
|
|
|
|
|
|
def test_serialize_works_when_dst_lacks_set_used_size() -> None:
|
|
"""Backward compatibility: a duck-typed destination that does NOT
|
|
implement ``set_used_size`` (e.g. a bare ``object()`` in older
|
|
tests) must still produce a successful task; the processor guards
|
|
the narrowing call with ``hasattr``."""
|
|
processor = AsyncSerdeProcessor(_FakeSerializer(), _FakeDeserializer())
|
|
try:
|
|
task_id = processor.submit_serialize([object()], [object()]) # type: ignore[list-item]
|
|
assert _wait_for_fd(processor.get_serialize_event_fd()), "fd never signaled"
|
|
# Success despite dst lacking set_used_size.
|
|
assert processor.query_serialize_result(task_id) is True
|
|
finally:
|
|
processor.close()
|
|
|
|
|
|
def test_serialize_skips_narrowing_when_serializer_returns_non_int() -> None:
|
|
"""A serializer that does not honor the ``serialize() -> int``
|
|
contract and returns ``None`` must not crash the processor. The
|
|
narrowing call is skipped (``isinstance(n, int)`` guard) rather than
|
|
passing ``None`` into ``set_used_size``, where the range check would
|
|
raise ``TypeError: '<' not supported between NoneType and int``."""
|
|
|
|
class _NoneReturningSerializer(Serializer):
|
|
def serialize(self, src, dst): # type: ignore[no-untyped-def]
|
|
return None # type: ignore[return-value]
|
|
|
|
def estimate_serialized_size(self, layout_desc: MemoryLayoutDesc) -> int:
|
|
return 1
|
|
|
|
processor = AsyncSerdeProcessor(_NoneReturningSerializer(), _FakeDeserializer())
|
|
try:
|
|
dst = _NarrowableDst()
|
|
task_id = processor.submit_serialize([object()], [dst]) # type: ignore[list-item]
|
|
assert _wait_for_fd(processor.get_serialize_event_fd()), "fd never signaled"
|
|
# Task still succeeds; narrowing was skipped, not attempted.
|
|
assert processor.query_serialize_result(task_id) is True
|
|
assert dst.used_sizes == [], "must not call set_used_size with a non-int"
|
|
finally:
|
|
processor.close()
|