94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
"""Coverage for bounded, ordered prefill transfer."""
|
|
|
|
from threading import Event
|
|
|
|
import pytest
|
|
|
|
from meshnet_gateway.prefill_backpressure import (
|
|
BoundedPrefillSender,
|
|
DEFAULT_PREFILL_CHUNK_TOKENS,
|
|
DEFAULT_PREFILL_MAX_IN_FLIGHT,
|
|
PrefillTransferLimits,
|
|
)
|
|
from meshnet_gateway.server import _BinaryActivation, _post_binary_forward
|
|
|
|
|
|
def test_limits_have_safe_defaults_and_keep_legacy_chunk_setting(monkeypatch):
|
|
monkeypatch.delenv("MESHNET_PREFILL_CHUNK_TOKENS", raising=False)
|
|
monkeypatch.delenv("MESHNET_PREFILL_MAX_IN_FLIGHT", raising=False)
|
|
monkeypatch.setenv("MESHNET_CHUNK_TOKENS", "64")
|
|
|
|
limits = PrefillTransferLimits.from_env()
|
|
|
|
assert limits.chunk_tokens == 64
|
|
assert limits.max_in_flight == DEFAULT_PREFILL_MAX_IN_FLIGHT
|
|
assert PrefillTransferLimits().chunk_tokens == DEFAULT_PREFILL_CHUNK_TOKENS
|
|
assert limits.max_buffered_bytes == limits.max_chunk_bytes
|
|
|
|
|
|
def test_slow_consumer_applies_backpressure_preserves_order_and_bounds_bytes():
|
|
sender = BoundedPrefillSender(
|
|
PrefillTransferLimits(chunk_tokens=2, max_in_flight=4, max_chunk_bytes=8)
|
|
)
|
|
sent: list[int] = []
|
|
produced: list[int] = []
|
|
|
|
def chunks():
|
|
for index in range(3):
|
|
produced.append(index)
|
|
yield bytes([index]) * 8
|
|
|
|
def forward(chunk: bytes) -> int:
|
|
sent.append(chunk[0])
|
|
# The next item cannot have been constructed while this consumer waits.
|
|
assert produced[-1] == chunk[0]
|
|
assert len(produced) == chunk[0] + 1
|
|
assert sender.buffered_bytes == 8
|
|
assert sender.in_flight == 1
|
|
return chunk[0]
|
|
|
|
assert sender.send(chunks(), body_size=len, forward=forward) == [0, 1, 2]
|
|
assert sent == [0, 1, 2]
|
|
assert sender.peak_buffered_bytes <= sender.limits.max_buffered_bytes
|
|
assert sender.peak_in_flight == 1
|
|
assert sender.buffered_bytes == sender.in_flight == 0
|
|
|
|
|
|
def test_failure_and_cancellation_release_owned_buffers():
|
|
sender = BoundedPrefillSender(PrefillTransferLimits())
|
|
with pytest.raises(RuntimeError, match="route lost"):
|
|
sender.send([b"one", b"two"], body_size=len, forward=lambda _: (_ for _ in ()).throw(RuntimeError("route lost")))
|
|
assert sender.closed
|
|
assert sender.buffered_bytes == sender.in_flight == 0
|
|
|
|
cancelled = Event()
|
|
cancelled.set()
|
|
sender = BoundedPrefillSender(PrefillTransferLimits())
|
|
assert sender.send([b"never"], body_size=len, forward=lambda _: None, cancelled=cancelled) == []
|
|
assert sender.buffered_bytes == sender.in_flight == 0
|
|
|
|
|
|
def test_legacy_single_chunk_peer_response_uses_outgoing_metadata(monkeypatch):
|
|
class Response:
|
|
headers = {"Content-Type": "application/octet-stream"}
|
|
|
|
def read(self):
|
|
return b"\0" * (1 * 1 * 64 * 2)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_):
|
|
return False
|
|
|
|
monkeypatch.setattr("urllib.request.urlopen", lambda *_args, **_kwargs: Response())
|
|
activation = _BinaryActivation(
|
|
body=b"\0" * (1 * 1 * 64 * 2), shape=[1, 1, 64], dtype="bfloat16",
|
|
session="legacy-session", chunk_index=0, chunk_total=1, encoding=None, headers={},
|
|
)
|
|
|
|
response = _post_binary_forward("http://legacy/forward", activation, hop_index=0, timeout=1.0)
|
|
|
|
assert response.session == activation.session
|
|
assert response.chunk_index == response.chunk_total - 1 == 0
|