popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: distributed GGUF shard load integration test with TinyLlama 1.1B

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-14 13:01:51 +03:00
parent ec36290863
commit 31065c0e12
2 changed files with 301 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

233
tests/test_gguf_distributed_load.py Normal file
View File

@@ -0,0 +1,233 @@
"""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()