463 lines
15 KiB
Python
463 lines
15 KiB
Python
"""US-012 tests for the real PyTorch node backend."""
|
|
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import threading
|
|
import time
|
|
import types
|
|
import urllib.request
|
|
|
|
import pytest
|
|
|
|
from meshnet_node.model_backend import (
|
|
InsufficientVRAMError,
|
|
TensorPayload,
|
|
_call_layer,
|
|
_decoder_attention_mask,
|
|
_int_tensor_header,
|
|
build_quantization_config,
|
|
validate_quantization,
|
|
)
|
|
from meshnet_node.torch_server import TorchNodeServer
|
|
|
|
|
|
class _FakeBackend:
|
|
model_id = "fake-model"
|
|
total_layers = 12
|
|
is_head = True
|
|
is_tail = False
|
|
|
|
def encode_prompt(self, prompt: str) -> TensorPayload:
|
|
assert prompt == "The capital of France is"
|
|
return TensorPayload(
|
|
body=b"\x00" * (1 * 6 * 8 * 2),
|
|
shape=[1, 6, 8],
|
|
attention_mask_header=None,
|
|
position_ids_header=None,
|
|
)
|
|
|
|
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
|
|
assert shape == [1, 6, 8]
|
|
return TensorPayload(
|
|
body=body,
|
|
shape=shape,
|
|
attention_mask_header=attention_mask_header,
|
|
position_ids_header=position_ids_header,
|
|
)
|
|
|
|
|
|
class _FakeTailBackend(_FakeBackend):
|
|
is_head = False
|
|
is_tail = True
|
|
|
|
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
|
|
assert len(body) == 1 * 6 * 8 * 2
|
|
return " Paris"
|
|
|
|
|
|
class _FakeFullBackend(_FakeBackend):
|
|
is_head = True
|
|
is_tail = True
|
|
|
|
def generate_text(
|
|
self,
|
|
messages: list[dict],
|
|
max_new_tokens: int = 16,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
) -> str:
|
|
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
|
|
assert max_new_tokens == 7
|
|
assert temperature == 1.0
|
|
assert top_p == 1.0
|
|
return "56"
|
|
|
|
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
|
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
|
|
return 8
|
|
|
|
def count_text_tokens(self, text: str) -> int:
|
|
assert text == "56"
|
|
return 1
|
|
|
|
|
|
class _FakeChatTokenizer:
|
|
eos_token = ""
|
|
|
|
def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
|
|
assert add_generation_prompt is True
|
|
assert tokenize is False
|
|
return "debug prompt"
|
|
|
|
|
|
class _FakePipelineHeadBackend(_FakeBackend):
|
|
tokenizer = _FakeChatTokenizer()
|
|
|
|
def encode_prompt(self, prompt: str) -> TensorPayload:
|
|
assert prompt.startswith("debug prompt")
|
|
return TensorPayload(
|
|
body=b"\x00" * (1 * 6 * 8 * 2),
|
|
shape=[1, 6, 8],
|
|
attention_mask_header=None,
|
|
position_ids_header=None,
|
|
)
|
|
|
|
|
|
class _FakePipelineTailBackend(_FakeTailBackend):
|
|
def __init__(self) -> None:
|
|
self.start_layers: list[int | None] = []
|
|
|
|
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
|
|
self.start_layers.append(start_layer)
|
|
assert len(body) == 1 * 6 * 8 * 2
|
|
return " token"
|
|
|
|
|
|
class _BlockingStreamingTailBackend(_FakeTailBackend):
|
|
def __init__(self, second_token_release: threading.Event) -> None:
|
|
self._release = second_token_release
|
|
self.calls = 0
|
|
|
|
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
return " first"
|
|
self._release.wait(timeout=3.0)
|
|
return " second"
|
|
|
|
|
|
def test_quantization_flag_validation():
|
|
assert validate_quantization("bfloat16") == "bfloat16"
|
|
assert validate_quantization("int8") == "int8"
|
|
assert validate_quantization("nf4") == "nf4"
|
|
with pytest.raises(ValueError, match="quantization"):
|
|
validate_quantization("float32")
|
|
|
|
|
|
def test_node_package_declares_torch_dependency():
|
|
pyproject = Path("packages/node/pyproject.toml").read_text(encoding="utf-8")
|
|
|
|
assert '"torch>=' in pyproject
|
|
|
|
|
|
def test_bitsandbytes_configs_are_created_lazily(monkeypatch):
|
|
calls = []
|
|
|
|
class FakeBitsAndBytesConfig:
|
|
def __init__(self, **kwargs):
|
|
calls.append(kwargs)
|
|
|
|
monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(bfloat16="bf16"))
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"transformers",
|
|
types.SimpleNamespace(BitsAndBytesConfig=FakeBitsAndBytesConfig),
|
|
)
|
|
|
|
assert build_quantization_config("bfloat16") is None
|
|
build_quantization_config("int8")
|
|
build_quantization_config("nf4")
|
|
|
|
assert calls == [
|
|
{"load_in_8bit": True},
|
|
{
|
|
"load_in_4bit": True,
|
|
"bnb_4bit_quant_type": "nf4",
|
|
"bnb_4bit_compute_dtype": "bf16",
|
|
},
|
|
]
|
|
|
|
|
|
def test_head_forward_accepts_text_prompt_and_returns_bfloat16_activations():
|
|
node = TorchNodeServer(backend=_FakeBackend())
|
|
port = node.start()
|
|
try:
|
|
payload = json.dumps({"prompt": "The capital of France is"}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{port}/forward",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
body = resp.read()
|
|
headers = {key.lower(): value for key, value in resp.headers.items()}
|
|
|
|
assert len(body) == 1 * 6 * 8 * 2
|
|
assert headers["x-meshnet-shape"] == "1,6,8"
|
|
assert headers["x-meshnet-dtype"] == "bfloat16"
|
|
assert headers["x-meshnet-wire"] == "2"
|
|
finally:
|
|
node.stop()
|
|
|
|
|
|
def test_tail_forward_returns_text_completion_from_binary_activations():
|
|
node = TorchNodeServer(backend=_FakeTailBackend())
|
|
port = node.start()
|
|
try:
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{port}/forward",
|
|
data=b"\x00" * (1 * 6 * 8 * 2),
|
|
headers={
|
|
"Content-Type": "application/octet-stream",
|
|
"X-Meshnet-Shape": "1,6,8",
|
|
"X-Meshnet-Dtype": "bfloat16",
|
|
"X-Meshnet-Session": "session-1",
|
|
"X-Meshnet-Chunk-Index": "0",
|
|
"X-Meshnet-Chunk-Total": "1",
|
|
"X-Meshnet-Hop-Index": "1",
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
body = json.loads(resp.read())
|
|
|
|
assert body == {"text": " Paris"}
|
|
assert node.received_activations
|
|
assert node.forward_chunk_count == 1
|
|
finally:
|
|
node.stop()
|
|
|
|
|
|
def test_full_model_chat_completion_uses_generation_not_single_token_decode():
|
|
node = TorchNodeServer(backend=_FakeFullBackend())
|
|
port = node.start()
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "What is 7 times 8?"}],
|
|
"max_tokens": 7,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
body = json.loads(resp.read())
|
|
|
|
assert body["choices"][0]["message"]["content"] == "56"
|
|
assert body["usage"] == {"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9}
|
|
finally:
|
|
node.stop()
|
|
|
|
|
|
def test_pipeline_hop_logs_are_suppressed_without_debug(capsys):
|
|
tail_backend = _FakePipelineTailBackend()
|
|
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
|
|
tail = TorchNodeServer(backend=tail_backend)
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"max_tokens": 1,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Route": json.dumps([
|
|
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
|
]),
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
body = json.loads(resp.read())
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
out = capsys.readouterr().out
|
|
assert body["choices"][0]["message"]["content"] == " token"
|
|
assert tail_backend.start_layers == [22]
|
|
assert "pipeline hop 0:" not in out
|
|
assert "pipeline hop 0 returned text" not in out
|
|
|
|
|
|
def test_pipeline_hop_logs_are_enabled_with_debug(capsys):
|
|
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True, debug=True)
|
|
tail = TorchNodeServer(backend=_FakePipelineTailBackend())
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"max_tokens": 1,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Route": json.dumps([
|
|
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
|
]),
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
json.loads(resp.read())
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
out = capsys.readouterr().out
|
|
assert f" [node] pipeline hop 0: http://127.0.0.1:{tail_port} start_layer=22" in out
|
|
assert " [node] pipeline hop 0 returned text=' token'" in out
|
|
|
|
|
|
def test_split_shard_chat_streams_each_generated_token_incrementally():
|
|
release_second = threading.Event()
|
|
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
|
|
tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
response = None
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"stream": True,
|
|
"max_tokens": 2,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Route": json.dumps([
|
|
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
|
]),
|
|
},
|
|
method="POST",
|
|
)
|
|
response = urllib.request.urlopen(req, timeout=5)
|
|
|
|
first_token_line = ""
|
|
deadline = time.time() + 2.0
|
|
while time.time() < deadline:
|
|
line = response.readline().decode()
|
|
if '"content": " first"' in line:
|
|
first_token_line = line
|
|
break
|
|
|
|
assert first_token_line
|
|
assert not release_second.is_set()
|
|
release_second.set()
|
|
rest = response.read().decode()
|
|
finally:
|
|
release_second.set()
|
|
if response is not None:
|
|
response.close()
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
assert '"content": " second"' in rest
|
|
assert "data: [DONE]" in rest
|
|
|
|
|
|
def test_int_tensor_header_serializes_torch_tensors():
|
|
torch = pytest.importorskip("torch")
|
|
|
|
header = _int_tensor_header(torch.tensor([[1, 2, 3]], dtype=torch.long))
|
|
|
|
assert header.startswith("1,3:")
|
|
|
|
|
|
def test_decoder_attention_mask_is_causal_float_mask():
|
|
torch = pytest.importorskip("torch")
|
|
|
|
hidden_states = torch.zeros((1, 3, 8), dtype=torch.bfloat16)
|
|
mask = _decoder_attention_mask(torch.ones((1, 3), dtype=torch.long), hidden_states, torch)
|
|
|
|
assert mask.shape == (1, 1, 3, 3)
|
|
assert mask.dtype == torch.bfloat16
|
|
assert mask[0, 0, 0, 1] < 0
|
|
assert mask[0, 0, 2, 0] == 0
|
|
|
|
|
|
def test_call_layer_passes_rotary_position_embeddings():
|
|
class NeedsPositionEmbeddings:
|
|
def __call__(self, hidden_states, **kwargs):
|
|
assert kwargs["position_embeddings"] == "rotary"
|
|
return hidden_states
|
|
|
|
assert _call_layer(
|
|
NeedsPositionEmbeddings(),
|
|
"hidden",
|
|
attention_mask=None,
|
|
position_ids="positions",
|
|
position_embeddings="rotary",
|
|
) == "hidden"
|
|
|
|
|
|
@pytest.mark.integration
|
|
def test_two_node_gpt2_completion_is_deterministic():
|
|
if os.environ.get("CI"):
|
|
pytest.skip("GPT-2 integration test is skipped in CI")
|
|
torch = pytest.importorskip("torch")
|
|
pytest.importorskip("transformers")
|
|
pytest.importorskip("safetensors")
|
|
pytest.importorskip("accelerate")
|
|
pytest.importorskip("bitsandbytes")
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("GPT-2 integration test requires a CUDA GPU")
|
|
|
|
head = TorchNodeServer(
|
|
model_id="openai-community/gpt2",
|
|
shard_start=0,
|
|
shard_end=6,
|
|
quantization="bfloat16",
|
|
)
|
|
tail = TorchNodeServer(
|
|
model_id="openai-community/gpt2",
|
|
shard_start=6,
|
|
shard_end=12,
|
|
quantization="bfloat16",
|
|
)
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
prompt_req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/forward",
|
|
data=json.dumps({"prompt": "The capital of France is"}).encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(prompt_req, timeout=60) as resp:
|
|
activation = resp.read()
|
|
head_headers = resp.headers
|
|
|
|
tail_req = urllib.request.Request(
|
|
f"http://127.0.0.1:{tail_port}/forward",
|
|
data=activation,
|
|
headers={
|
|
"Content-Type": "application/octet-stream",
|
|
"X-Meshnet-Shape": head_headers["X-Meshnet-Shape"],
|
|
"X-Meshnet-Dtype": head_headers["X-Meshnet-Dtype"],
|
|
"X-Meshnet-Session": "gpt2-session",
|
|
"X-Meshnet-Chunk-Index": "0",
|
|
"X-Meshnet-Chunk-Total": "1",
|
|
"X-Meshnet-Hop-Index": "1",
|
|
"X-Meshnet-Attn-Mask": head_headers["X-Meshnet-Attn-Mask"],
|
|
"X-Meshnet-Position-Ids": head_headers["X-Meshnet-Position-Ids"],
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(tail_req, timeout=60) as resp:
|
|
body = json.loads(resp.read())
|
|
|
|
assert body["text"].strip()
|
|
assert body["text"] == " Paris"
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|