116 lines
4.2 KiB
Python
116 lines
4.2 KiB
Python
"""Integration tests for two-node shard pipelines with activation tensors flowing A to B."""
|
|
|
|
import json
|
|
import urllib.request
|
|
|
|
from meshnet_node.server import StubNodeServer
|
|
from meshnet_gateway.server import GatewayServer
|
|
|
|
|
|
def test_two_node_activation_pipeline():
|
|
# Node A owns layers 0-15 (first shard, not last — returns activations).
|
|
node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
|
|
port_a = node_a.start()
|
|
|
|
# Node B owns layers 16-31 (last shard — returns text after receiving activations).
|
|
node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
|
|
port_b = node_b.start()
|
|
|
|
# Gateway hardcodes the two-node inference route for this story.
|
|
gateway = GatewayServer(
|
|
inference_route=[
|
|
f"http://127.0.0.1:{port_a}",
|
|
f"http://127.0.0.1:{port_b}",
|
|
]
|
|
)
|
|
gateway_port = gateway.start()
|
|
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "stub-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
}).encode()
|
|
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
|
|
with urllib.request.urlopen(req) as resp:
|
|
assert resp.status == 200
|
|
body = json.loads(resp.read())
|
|
|
|
# Gateway must assemble a valid OpenAI-format response.
|
|
assert body["id"] == "chatcmpl-stub"
|
|
assert body["object"] == "chat.completion"
|
|
assert body["model"] == "stub-model"
|
|
assert len(body["choices"]) == 1
|
|
choice = body["choices"][0]
|
|
assert choice["message"]["role"] == "assistant"
|
|
assert choice["message"]["content"] == "stub response to: hello"
|
|
assert choice["finish_reason"] == "stop"
|
|
|
|
# Node B must have received an activation tensor from node A — this is
|
|
# the observable proof that the A → B activation hop occurred.
|
|
assert node_b.received_activations, "node B did not receive activation tensors from node A"
|
|
|
|
# Node A must NOT have received activations (it is the first shard).
|
|
assert not node_a.received_activations, "node A should not receive activations — it is first in the route"
|
|
finally:
|
|
gateway.stop()
|
|
node_a.stop()
|
|
node_b.stop()
|
|
|
|
|
|
def test_binary_activation_pipeline_chunks_512_token_prefill():
|
|
node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
|
|
port_a = node_a.start()
|
|
node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
|
|
port_b = node_b.start()
|
|
gateway = GatewayServer(
|
|
inference_route=[
|
|
f"http://127.0.0.1:{port_a}",
|
|
f"http://127.0.0.1:{port_b}",
|
|
]
|
|
)
|
|
gateway_port = gateway.start()
|
|
|
|
try:
|
|
prompt = " ".join(f"tok{i}" for i in range(512))
|
|
payload = json.dumps({
|
|
"model": "stub-model",
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
|
|
with urllib.request.urlopen(req) as resp:
|
|
assert resp.status == 200
|
|
body = json.loads(resp.read())
|
|
|
|
assert body["choices"][0]["message"]["content"] == f"stub response to: {prompt}"
|
|
assert node_a.forward_chunk_count == 4
|
|
assert node_b.forward_chunk_count == 4
|
|
assert not node_a.received_activations
|
|
assert node_b.received_activations
|
|
|
|
chunk_responses = gateway.last_binary_chunk_responses
|
|
assert len(chunk_responses) == 4
|
|
for index, chunk in enumerate(chunk_responses):
|
|
assert chunk.shape == [1, 128, 64]
|
|
assert chunk.dtype == "bfloat16"
|
|
assert chunk.chunk_index == index
|
|
assert chunk.chunk_total == 4
|
|
assert chunk.headers["x-meshnet-wire"] == "2"
|
|
assert len(chunk.body) == 1 * 128 * 64 * 2
|
|
finally:
|
|
gateway.stop()
|
|
node_a.stop()
|
|
node_b.stop()
|