233 lines
8.0 KiB
Python
233 lines
8.0 KiB
Python
"""Distributed GGUF shard load integration test.
|
|
|
|
Downloads a small dense-Llama GGUF (TinyLlama 1.1B Q4_K_M ~670 MB),
|
|
loads it in shard ranges via the meshnet-range-loader C wrapper, registers
|
|
each shard with a live TrackerServer, and verifies routing, range reporting,
|
|
and memory scaling.
|
|
|
|
Set MESHNET_ENABLE_REAL_INFERENCE_TESTS=1 to run.
|
|
|
|
Evidence class: real-model integration. Downloads ~670 MB on first run.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import subprocess
|
|
import sys
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
import pytest
|
|
|
|
ROOT = pathlib.Path(__file__).resolve().parent.parent
|
|
sys.path[:0] = [
|
|
str(ROOT / "packages" / "tracker"),
|
|
str(ROOT / "packages" / "node"),
|
|
str(ROOT / "packages" / "contracts"),
|
|
]
|
|
|
|
from meshnet_tracker.server import TrackerServer # noqa: E402 — sys.path prepended above
|
|
|
|
# Only run when explicitly enabled
|
|
pytestmark = pytest.mark.skipif(
|
|
"MESHNET_ENABLE_REAL_INFERENCE_TESTS" not in os.environ,
|
|
reason="set MESHNET_ENABLE_REAL_INFERENCE_TESTS=1 to download and load a real GGUF",
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
HF_REPO = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF"
|
|
GGUF_FILE = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
|
|
MODEL_URL = f"https://huggingface.co/{HF_REPO}/resolve/main/{GGUF_FILE}"
|
|
EXPECTED_LAYERS = 22
|
|
GGUF_FLAVOR = GGUF_FILE.replace(".gguf", "")
|
|
|
|
CACHE_DIR = ROOT / ".cache" / "gguf-models"
|
|
LOADER = ROOT / "build" / "dgr-004-final" / "build" / "bin" / "meshnet-range-loader"
|
|
LLAMA_LIB_DIR = LOADER.parent
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Loading helper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_shard(gguf_path: str, start: int, end: int) -> dict:
|
|
"""Load a shard range of the GGUF via the C wrapper and return JSON report."""
|
|
env = os.environ.copy()
|
|
env["LD_LIBRARY_PATH"] = str(LLAMA_LIB_DIR)
|
|
|
|
result = subprocess.run(
|
|
[str(LOADER), gguf_path, str(start), str(end)],
|
|
capture_output=True, text=True, timeout=120, env=env,
|
|
)
|
|
if result.returncode != 0:
|
|
# Parse JSON from stdout even on error if it exists
|
|
if result.stdout.strip():
|
|
try:
|
|
return json.loads(result.stdout)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
raise RuntimeError(
|
|
f"meshnet-range-loader [{start}, {end}) failed (exit {result.returncode}):\n"
|
|
f"stderr: {result.stderr[:500]}"
|
|
)
|
|
|
|
return json.loads(result.stdout)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Download and cache
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _ensure_model() -> pathlib.Path:
|
|
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
|
model_path = CACHE_DIR / GGUF_FILE
|
|
|
|
if model_path.exists() and model_path.stat().st_size > 600 * 1024 * 1024:
|
|
return model_path
|
|
|
|
print(f"Downloading {MODEL_URL} (~670 MB)...", file=sys.stderr)
|
|
urllib.request.urlretrieve(MODEL_URL, model_path)
|
|
actual_mb = model_path.stat().st_size / (1024 * 1024)
|
|
print(f"Downloaded {GGUF_FLAVOR}: {actual_mb:.0f} MB", file=sys.stderr)
|
|
return model_path
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tracker helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _post_json(url: str, data: dict) -> dict:
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=json.dumps(data).encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=10) as r:
|
|
return json.loads(r.read())
|
|
|
|
|
|
def _get_json(url: str) -> dict:
|
|
with urllib.request.urlopen(url, timeout=10) as r:
|
|
return json.loads(r.read())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_whole_model_load_and_report():
|
|
"""Load the full TinyLlama GGUF and verify metadata."""
|
|
gguf = _ensure_model()
|
|
report = load_shard(str(gguf), 0, EXPECTED_LAYERS)
|
|
assert report["ok"] is True
|
|
assert report["start_layer"] == 0
|
|
assert report["end_layer"] == EXPECTED_LAYERS
|
|
assert report["mapped_bytes"] > 0
|
|
assert report["resident_bytes"] >= report["mapped_bytes"]
|
|
|
|
|
|
def test_head_shard_is_less_than_full_model():
|
|
"""A head-only shard maps fewer bytes than the full model."""
|
|
gguf = _ensure_model()
|
|
head = load_shard(str(gguf), 0, 8)
|
|
full = load_shard(str(gguf), 0, EXPECTED_LAYERS)
|
|
assert head["mapped_bytes"] <= full["mapped_bytes"]
|
|
|
|
|
|
def test_tail_shard_maps_fewer_bytes_than_middle():
|
|
"""Fewer layers = fewer bytes (tail has 6 layers, middle has 8)."""
|
|
gguf = _ensure_model()
|
|
middle = load_shard(str(gguf), 8, 16)
|
|
tail = load_shard(str(gguf), 16, EXPECTED_LAYERS)
|
|
# Middle (8 layers) must map more than tail (6 layers)
|
|
assert middle["mapped_bytes"] > tail["mapped_bytes"]
|
|
|
|
|
|
def test_memory_scales_with_owned_range():
|
|
"""More layers = more resident bytes."""
|
|
gguf = _ensure_model()
|
|
small = load_shard(str(gguf), 0, 4)
|
|
large = load_shard(str(gguf), 0, 12)
|
|
assert large["mapped_bytes"] > small["mapped_bytes"]
|
|
assert large["resident_bytes"] > small["resident_bytes"]
|
|
|
|
|
|
def test_invalid_range_rejected():
|
|
"""Loading a range outside GGUF layer bounds fails closed."""
|
|
gguf = _ensure_model()
|
|
env = os.environ.copy()
|
|
env["LD_LIBRARY_PATH"] = str(LLAMA_LIB_DIR)
|
|
result = subprocess.run(
|
|
[str(LOADER), str(gguf), str(EXPECTED_LAYERS + 1), str(EXPECTED_LAYERS + 5)],
|
|
capture_output=True, text=True, timeout=30, env=env,
|
|
)
|
|
# Should fail with exit code 1 and an error message
|
|
assert result.returncode != 0
|
|
assert "error" in result.stderr.lower() or "ERROR" in result.stderr
|
|
|
|
|
|
def test_tracker_registers_shard_nodes():
|
|
"""Start a tracker, register multiple shard nodes, verify the console."""
|
|
gguf = _ensure_model()
|
|
|
|
tracker = TrackerServer(heartbeat_timeout=60.0)
|
|
tracker_port = tracker.start()
|
|
|
|
shards = [
|
|
("head", 0, 8),
|
|
("middle", 8, 16),
|
|
("tail", 16, EXPECTED_LAYERS),
|
|
]
|
|
registered_ids = []
|
|
try:
|
|
for name_tag, start, end in shards:
|
|
report = load_shard(str(gguf), start, end)
|
|
node_id = f"{GGUF_FLAVOR}-{name_tag}"
|
|
registered_ids.append(node_id)
|
|
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{
|
|
"node_id": node_id,
|
|
"endpoint": f"http://localhost:{10000 + start}",
|
|
"model": GGUF_FLAVOR,
|
|
"num_layers": report["end_layer"] - report["start_layer"],
|
|
"shard_start": report["start_layer"],
|
|
"shard_end": report["end_layer"],
|
|
"hardware_profile": {
|
|
"mapped_bytes": report["mapped_bytes"],
|
|
"resident_bytes": report["resident_bytes"],
|
|
},
|
|
"score": 1.0,
|
|
},
|
|
)
|
|
|
|
# Verify console shows all three registered nodes
|
|
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
|
|
registered_eps = {
|
|
f"http://localhost:{10000 + start}" for _, start, _ in shards
|
|
}
|
|
found_eps = set()
|
|
for event in console.get("events", []):
|
|
if event.get("message") == "node registered":
|
|
ep = event.get("fields", {}).get("endpoint", "")
|
|
if ep in registered_eps:
|
|
found_eps.add(ep)
|
|
for _, start, _ in shards:
|
|
expected_ep = f"http://localhost:{10000 + start}"
|
|
assert expected_ep in found_eps, \
|
|
f"endpoint {expected_ep} not found in registration events"
|
|
|
|
finally:
|
|
tracker.stop() |