605 lines
22 KiB
Python
605 lines
22 KiB
Python
"""US-003 tests: tracker node registration and route selection."""
|
|
|
|
import http.server
|
|
import json
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
from meshnet_gateway.server import GatewayServer, _banned_route_wallet
|
|
from meshnet_node.server import StubNodeServer
|
|
from meshnet_contracts import LocalSolanaContracts
|
|
from meshnet_tracker.server import TrackerServer, _registration_ban_error
|
|
|
|
|
|
def _post_json(url: str, payload: dict) -> dict:
|
|
data = json.dumps(payload).encode()
|
|
req = urllib.request.Request(
|
|
url, data=data, headers={"Content-Type": "application/json"}, method="POST"
|
|
)
|
|
with urllib.request.urlopen(req) as r:
|
|
return json.loads(r.read())
|
|
|
|
|
|
def _get_json(url: str) -> dict:
|
|
with urllib.request.urlopen(url) as r:
|
|
return json.loads(r.read())
|
|
|
|
|
|
def test_tracker_node_registration():
|
|
"""A node can register with the tracker and receives a node_id."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
resp = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{
|
|
"endpoint": "http://127.0.0.1:9000",
|
|
"shard_start": 0,
|
|
"shard_end": 15,
|
|
"hardware_profile": {"gpu": "rtx3090"},
|
|
"score": 1.0,
|
|
},
|
|
)
|
|
assert "node_id" in resp
|
|
assert isinstance(resp["node_id"], str)
|
|
assert len(resp["node_id"]) > 0
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_route_selection():
|
|
"""Tracker returns ordered route when nodes collectively cover all layers."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
route_resp = _get_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model"
|
|
)
|
|
assert "route" in route_resp
|
|
assert route_resp["route"] == [
|
|
"http://127.0.0.1:9001",
|
|
"http://127.0.0.1:9002",
|
|
]
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_route_error_no_coverage():
|
|
"""Tracker returns 503 when registered nodes do not cover all required layers."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
# Only half the layers covered — route must fail.
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 but request succeeded")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
body = json.loads(exc.read())
|
|
assert "error" in body
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_route_error_no_nodes():
|
|
"""Tracker returns 503 with clear error when the registry is empty."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 but request succeeded")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
body = json.loads(exc.read())
|
|
assert "error" in body
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_heartbeat_updates_node():
|
|
"""Sending a heartbeat for a registered node succeeds."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
reg = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
node_id = reg["node_id"]
|
|
|
|
hb_resp = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {}
|
|
)
|
|
assert hb_resp == {}
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_heartbeat_expiry():
|
|
"""Nodes that miss their heartbeat window are excluded from routes."""
|
|
tracker = TrackerServer(heartbeat_timeout=0.05) # 50 ms
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
# Immediately after registration the node is alive.
|
|
route_resp = _get_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model"
|
|
)
|
|
assert route_resp["route"] == ["http://127.0.0.1:9001"]
|
|
|
|
# Let the heartbeat window expire.
|
|
time.sleep(0.15)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 after heartbeat expiry")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_heartbeat_expiry_removes_node_from_registry():
|
|
"""Expired nodes are removed, not merely hidden from route responses."""
|
|
tracker = TrackerServer(heartbeat_timeout=0.05)
|
|
tracker_port = tracker.start()
|
|
try:
|
|
reg = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
node_id = reg["node_id"]
|
|
time.sleep(0.15)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 after heartbeat expiry")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
|
|
try:
|
|
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {})
|
|
raise AssertionError("Expected 404 for removed node")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 404
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_route_rejects_non_extending_overlap():
|
|
"""Overlapping shards that do not extend coverage cannot form a complete route."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 10,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "shard_start": 0, "shard_end": 5,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 for incomplete overlapping route")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_registration_rejects_invalid_payload():
|
|
"""Registration errors return a defined JSON 400 response."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "not-a-url", "shard_start": 10, "shard_end": 0,
|
|
"hardware_profile": []},
|
|
)
|
|
raise AssertionError("Expected 400 for invalid registration")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 400
|
|
body = json.loads(exc.read())
|
|
assert "error" in body
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_routes_only_nodes_for_requested_model():
|
|
"""A node registered for one model cannot satisfy another model's route."""
|
|
tracker = TrackerServer(model_presets={
|
|
"model-a": {"layers_start": 0, "layers_end": 31},
|
|
"model-b": {"layers_start": 0, "layers_end": 31},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
|
"model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=model-b")
|
|
raise AssertionError("Expected model-b route to ignore model-a nodes")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_assignment_coverage_is_model_scoped():
|
|
"""Shard assignment gaps are computed only from nodes serving the same model."""
|
|
tracker = TrackerServer(model_presets={
|
|
"model-a": {"layers_start": 0, "layers_end": 31},
|
|
"model-b": {"layers_start": 0, "layers_end": 31},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
|
"model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
|
|
assignment = _get_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/assign?model=model-b&device=cpu")
|
|
assert assignment["shard_start"] == 0
|
|
assert assignment["shard_end"] == 15
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_excludes_banned_wallets_from_routes():
|
|
"""Tracker refuses route candidates whose wallet is banned on-chain."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.submit_stake("wallet-a", 1_000)
|
|
contracts.registry.submit_stake("wallet-b", 1_000)
|
|
contracts.registry.ban_wallet("wallet-b")
|
|
tracker = TrackerServer(contracts=contracts, minimum_stake=100)
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
|
|
)
|
|
raise AssertionError("Expected 403 for banned wallet registration")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 403
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 when only banned wallet covers a layer")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
body = json.loads(exc.read())
|
|
assert "layer 16" in body["error"]
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_rejects_banned_wallet_registration():
|
|
"""A banned wallet cannot register with the tracker."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.submit_stake("wallet-b", 1_000)
|
|
contracts.registry.ban_wallet("wallet-b")
|
|
tracker = TrackerServer(contracts=contracts, minimum_stake=100)
|
|
tracker_port = tracker.start()
|
|
try:
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "shard_start": 0, "shard_end": 31,
|
|
"hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
|
|
)
|
|
raise AssertionError("Expected 403 for banned wallet registration")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 403
|
|
body = json.loads(exc.read())
|
|
assert body == {"error": "wallet is banned"}
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_ban_registration_guard_reads_contract_state():
|
|
"""The tracker registration guard reads ban status from the contract facade."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.ban_wallet("wallet-b")
|
|
|
|
assert _registration_ban_error(contracts, "wallet-b") == "wallet is banned"
|
|
assert _registration_ban_error(contracts, "wallet-a") is None
|
|
assert _registration_ban_error(None, "wallet-b") is None
|
|
|
|
|
|
def test_gateway_route_guard_rejects_banned_wallet_metadata():
|
|
"""Gateway rejects tracker routes that include a wallet now banned on-chain."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.ban_wallet("wallet-b")
|
|
|
|
banned_wallet = _banned_route_wallet(
|
|
contracts,
|
|
[{"endpoint": "http://node", "wallet_address": "wallet-b"}],
|
|
)
|
|
|
|
assert banned_wallet == "wallet-b"
|
|
|
|
|
|
def test_gateway_routes_even_when_compute_node_has_low_stake():
|
|
"""Routing is not based on a compute node's token balance."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.submit_stake("wallet-a", 1_000)
|
|
contracts.registry.submit_stake("wallet-b", 5)
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
|
|
|
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()
|
|
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
|
|
)
|
|
|
|
gateway = GatewayServer(
|
|
tracker_url=tracker_url,
|
|
contracts=contracts,
|
|
minimum_stake=100,
|
|
)
|
|
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", "Authorization": "Bearer api-key-a"},
|
|
method="POST",
|
|
)
|
|
|
|
with urllib.request.urlopen(req) as resp:
|
|
assert resp.status == 200
|
|
assert node_b.received_activations
|
|
finally:
|
|
gateway.stop()
|
|
node_a.stop()
|
|
node_b.stop()
|
|
tracker.stop()
|
|
|
|
|
|
def test_gateway_records_compute_attribution_after_inference_session():
|
|
"""Gateway records layer and token attribution for every routed node."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.submit_stake("wallet-a", 1_000)
|
|
contracts.registry.submit_stake("wallet-b", 1_000)
|
|
contracts.payment.fund_api_key("api-key-a", lamports=10_000)
|
|
tracker = TrackerServer(contracts=contracts, minimum_stake=100)
|
|
tracker_port = tracker.start()
|
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
|
|
|
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()
|
|
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
|
|
)
|
|
|
|
gateway = GatewayServer(
|
|
tracker_url=tracker_url,
|
|
contracts=contracts,
|
|
minimum_stake=100,
|
|
)
|
|
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", "Authorization": "Bearer api-key-a"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req) as resp:
|
|
assert resp.status == 200
|
|
|
|
records = contracts.payment.list_attributions()
|
|
assert [(r.api_key, r.node_wallet, r.layer_start, r.layer_end, r.tokens) for r in records] == [
|
|
("api-key-a", "wallet-a", 0, 15, 1),
|
|
("api-key-a", "wallet-b", 16, 31, 1),
|
|
]
|
|
finally:
|
|
gateway.stop()
|
|
node_a.stop()
|
|
node_b.stop()
|
|
tracker.stop()
|
|
|
|
|
|
def test_gateway_returns_json_503_when_tracker_unavailable():
|
|
"""Tracker-backed gateway failures are reported as JSON HTTP errors."""
|
|
gateway = GatewayServer(tracker_url="http://127.0.0.1:9")
|
|
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",
|
|
)
|
|
try:
|
|
urllib.request.urlopen(req, timeout=5)
|
|
raise AssertionError("Expected 503 for unavailable tracker")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
assert exc.headers["Content-Type"] == "application/json"
|
|
body = json.loads(exc.read())
|
|
assert body == {"error": "tracker unavailable"}
|
|
finally:
|
|
gateway.stop()
|
|
|
|
|
|
def test_gateway_returns_json_503_for_malformed_tracker_route():
|
|
"""Malformed tracker route payloads do not drop the gateway connection."""
|
|
|
|
class MalformedRouteHandler(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, fmt, *args):
|
|
pass
|
|
|
|
def do_GET(self):
|
|
body = json.dumps({"route": []}).encode()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
fake_tracker = http.server.HTTPServer(("127.0.0.1", 0), MalformedRouteHandler)
|
|
fake_thread = threading.Thread(target=fake_tracker.serve_forever, daemon=True)
|
|
fake_thread.start()
|
|
gateway = GatewayServer(tracker_url=f"http://127.0.0.1:{fake_tracker.server_address[1]}")
|
|
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",
|
|
)
|
|
try:
|
|
urllib.request.urlopen(req, timeout=5)
|
|
raise AssertionError("Expected 503 for malformed tracker route")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
assert exc.headers["Content-Type"] == "application/json"
|
|
body = json.loads(exc.read())
|
|
assert body == {"error": "tracker unavailable"}
|
|
finally:
|
|
gateway.stop()
|
|
fake_tracker.shutdown()
|
|
fake_tracker.server_close()
|
|
fake_thread.join(timeout=1)
|
|
|
|
|
|
def test_two_node_pipeline_via_tracker():
|
|
"""End-to-end: nodes register with tracker; gateway discovers route dynamically."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
|
|
|
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()
|
|
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
gateway = GatewayServer(tracker_url=tracker_url)
|
|
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())
|
|
|
|
assert body["id"] == "chatcmpl-stub"
|
|
assert body["object"] == "chat.completion"
|
|
assert body["choices"][0]["message"]["content"] == "stub response to: hello"
|
|
assert node_b.received_activations, "node B did not receive activations from node A"
|
|
assert not node_a.received_activations, "node A should not receive activations"
|
|
finally:
|
|
gateway.stop()
|
|
node_a.stop()
|
|
node_b.stop()
|
|
tracker.stop()
|