Files
neuron-tai/tests/test_real_distributed_inference.py

63 lines
2.2 KiB
Python

"""Environment-gated acceptance test against real registered model shards.
Run only after starting actual nodes (never synthetic HTTP handlers):
MESHNET_REAL_INFERENCE_URL=http://localhost:8080 \
MESHNET_REAL_INFERENCE_API_KEY=... \
MESHNET_REAL_INFERENCE_MODEL=Qwen/Qwen2.5-0.5B-Instruct \
MESHNET_REAL_INFERENCE_ROUTE=node-id-head,node-id-tail \
.venv-rocm/bin/python -m pytest tests/test_real_distributed_inference.py -m integration -v
"""
import json
import os
import time
import urllib.request
import pytest
BASE_URL = os.environ.get("MESHNET_REAL_INFERENCE_URL")
API_KEY = os.environ.get("MESHNET_REAL_INFERENCE_API_KEY")
MODEL = os.environ.get("MESHNET_REAL_INFERENCE_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
ROUTE = [node_id for node_id in os.environ.get("MESHNET_REAL_INFERENCE_ROUTE", "").split(",") if node_id]
pytestmark = pytest.mark.integration
def _required_real_stack() -> None:
if not BASE_URL or not API_KEY or len(ROUTE) < 2:
pytest.skip(
"set MESHNET_REAL_INFERENCE_URL, MESHNET_REAL_INFERENCE_API_KEY, and "
"a comma-separated MESHNET_REAL_INFERENCE_ROUTE with at least two real nodes"
)
def test_real_registered_shards_complete_a_pinned_request():
"""Acceptance test: tracker proxies an actual prompt through real registered shards."""
_required_real_stack()
assert BASE_URL is not None
assert API_KEY is not None
request = urllib.request.Request(
f"{BASE_URL.rstrip('/')}/v1/chat/completions",
data=json.dumps({
"model": MODEL,
"messages": [{"role": "user", "content": "What is 2 plus 2? Reply in one word."}],
"max_tokens": 8,
"temperature": 1.0,
"route": ROUTE,
}).encode(),
headers={"Content-Type": "application/json", "Authorization": f"Bearer {API_KEY}"},
method="POST",
)
started = time.monotonic()
with urllib.request.urlopen(request, timeout=300) as response:
payload = json.loads(response.read())
elapsed = time.monotonic() - started
content = payload["choices"][0]["message"]["content"].strip().lower()
assert content
assert "two" in content or "four" in content
assert elapsed > 0
assert payload["usage"]["total_tokens"] > 0