From 31065c0e120032d2bd1e9df71036e4720d671f97 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Tue, 14 Jul 2026 13:01:51 +0300 Subject: [PATCH] feat: distributed GGUF shard load integration test with TinyLlama 1.1B --- .../native/llama/meshnet-range-loader.cpp | 68 +++++ tests/test_gguf_distributed_load.py | 233 ++++++++++++++++++ 2 files changed, 301 insertions(+) create mode 100644 packages/node/native/llama/meshnet-range-loader.cpp create mode 100644 tests/test_gguf_distributed_load.py diff --git a/packages/node/native/llama/meshnet-range-loader.cpp b/packages/node/native/llama/meshnet-range-loader.cpp new file mode 100644 index 0000000..2161692 --- /dev/null +++ b/packages/node/native/llama/meshnet-range-loader.cpp @@ -0,0 +1,68 @@ +/** + * meshnet-range-loader — CLI wrapper that loads a GGUF shard range + * and prints the meshnet range report as JSON. + * + * Usage: meshnet-range-loader + * + * Build: g++ -std=c++17 -I/include -L/bin + * meshnet-range-loader.cpp -lllama -o meshnet-range-loader + * LD_LIBRARY_PATH=/bin ./meshnet-range-loader ... + */ + +#include "llama.h" + +#include +#include +#include + +int main(int argc, char **argv) { + if (argc != 4) { + fprintf(stderr, "Usage: %s \n", argv[0]); + return 1; + } + + const char *path = argv[1]; + int start = std::atoi(argv[2]); + int end = std::atoi(argv[3]); + + llama_backend_init(); + + llama_model_params params = llama_model_default_params(); + params.meshnet_owned_layer_start = start; + params.meshnet_owned_layer_end = end; + + llama_model *model = llama_model_load_from_file(path, params); + if (!model) { + fprintf(stderr, "ERROR model_load returned nullptr\n"); + llama_backend_free(); + return 1; + } + + llama_meshnet_range_report report; + bool ok = llama_model_meshnet_range_report(model, &report); + + int n_layer = llama_model_n_layer(model); + int n_embd = llama_model_n_embd(model); + uint64_t size = llama_model_size(model); + + if (ok) { + fprintf(stdout, + "{" + "\"ok\":true," + "\"start_layer\":%d,\"end_layer\":%d," + "\"mapped_bytes\":%lu,\"resident_bytes\":%lu," + "\"model_n_layer\":%d,\"model_n_embd\":%d,\"model_size\":%lu" + "}\n", + report.start_layer, report.end_layer, + (unsigned long)report.mapped_bytes, + (unsigned long)report.resident_bytes, + n_layer, n_embd, (unsigned long)size + ); + } else { + fprintf(stderr, "ERROR range report not available\n"); + } + + llama_model_free(model); + llama_backend_free(); + return ok ? 0 : 1; +} \ No newline at end of file diff --git a/tests/test_gguf_distributed_load.py b/tests/test_gguf_distributed_load.py new file mode 100644 index 0000000..35943f0 --- /dev/null +++ b/tests/test_gguf_distributed_load.py @@ -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() \ No newline at end of file