feat: add binary activation wire format
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
"""US-002 integration test: two-node shard pipeline with activation tensors flowing A → B."""
|
||||
"""Integration tests for two-node shard pipelines with activation tensors flowing A to B."""
|
||||
|
||||
import json
|
||||
import urllib.request
|
||||
@@ -62,3 +62,54 @@ def test_two_node_activation_pipeline():
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user