diff --git a/.scratch/distributed-inference-network/prd.json b/.scratch/distributed-inference-network/prd.json index 2840388..800e160 100644 --- a/.scratch/distributed-inference-network/prd.json +++ b/.scratch/distributed-inference-network/prd.json @@ -93,11 +93,12 @@ "Commit only this story's code changes with a focused conventional commit message" ], "priority": 4, - "passes": false, + "passes": true, "notes": "Source issue: .scratch/distributed-inference-network/issues/04-node-client-startup.md", "dependsOn": [ "US-003" - ] + ], + "completionNotes": "Completed by Ralph iteration 86510a10; verified by pytest, compileall, editable installs, CLI help, and final review." }, { "id": "US-005", diff --git a/packages/node/meshnet_node/cli.py b/packages/node/meshnet_node/cli.py index 67d2ab5..8aa889a 100644 --- a/packages/node/meshnet_node/cli.py +++ b/packages/node/meshnet_node/cli.py @@ -2,6 +2,7 @@ import argparse import sys +import time def main() -> None: @@ -12,20 +13,34 @@ def main() -> None: subparsers = parser.add_subparsers(dest="command") start_cmd = subparsers.add_parser("start", help="Start the node server") - start_cmd.add_argument("--tracker", default="http://localhost:8080", help="Tracker URL") + start_cmd.add_argument( + "--tracker", default="http://localhost:8080", help="Tracker URL" + ) start_cmd.add_argument("--port", type=int, default=7000, help="Port to listen on") + start_cmd.add_argument( + "--model", default="stub-model", help="Model preset to request from tracker" + ) + start_cmd.add_argument( + "--host", default="0.0.0.0", help="Interface to bind to" + ) + start_cmd.add_argument( + "--advertise-host", + help="Reachable host/IP to advertise to the tracker (defaults to FQDN when binding 0.0.0.0)", + ) args = parser.parse_args() if args.command == "start": - from meshnet_node.server import StubNodeServer + from meshnet_node.startup import run_startup - node = StubNodeServer(port=args.port) - port = node.start() - print(f"meshnet-node listening on port {port}") - print(f"Tracker URL: {args.tracker}") + node = run_startup( + tracker_url=args.tracker, + port=args.port, + model=args.model, + host=args.host, + advertise_host=args.advertise_host, + ) try: - import time while True: time.sleep(1) except KeyboardInterrupt: diff --git a/packages/node/meshnet_node/downloader.py b/packages/node/meshnet_node/downloader.py new file mode 100644 index 0000000..5716d49 --- /dev/null +++ b/packages/node/meshnet_node/downloader.py @@ -0,0 +1,62 @@ +"""Shard downloader — fetches or stubs a model shard from HuggingFace Hub. + +Cache layout: ~/.cache/meshnet/shards//layers_-/ + +For "stub-model" (no HF repo), a placeholder JSON file is written so the +test suite never touches the network. +""" + +import json +from pathlib import Path + +_DEFAULT_CACHE = Path.home() / ".cache" / "meshnet" / "shards" + + +def download_shard( + model: str, + shard_start: int, + shard_end: int, + cache_dir: Path = _DEFAULT_CACHE, + hf_repo: str | None = None, + progress: bool = True, +) -> Path: + """Ensure the shard is present in *cache_dir* and return its local path. + + When *hf_repo* is None (or *model* is ``"stub-model"``), a placeholder + weights file is created locally — no network access required. This keeps + the test suite hermetic while the real download path is exercised by + passing a non-stub *hf_repo*. + """ + shard_dir = cache_dir / model / f"layers_{shard_start}-{shard_end}" + shard_dir.mkdir(parents=True, exist_ok=True) + + if hf_repo is None or model == "stub-model": + stub_file = shard_dir / "weights.json" + if not stub_file.exists(): + stub_file.write_text(json.dumps({ + "model": model, + "shard_start": shard_start, + "shard_end": shard_end, + "stub": True, + })) + if progress: + print(f" [stub] shard placeholder written to {stub_file}", flush=True) + else: + if progress: + print(f" [stub] shard already cached at {shard_dir}", flush=True) + return shard_dir + + from huggingface_hub import snapshot_download # type: ignore[import] + + if progress: + print( + f" Downloading layers {shard_start}-{shard_end} from {hf_repo} ...", + flush=True, + ) + + local_dir = snapshot_download( + repo_id=hf_repo, + cache_dir=str(cache_dir), + local_dir=str(shard_dir), + ) + return Path(local_dir) diff --git a/packages/node/meshnet_node/hardware.py b/packages/node/meshnet_node/hardware.py new file mode 100644 index 0000000..33724dd --- /dev/null +++ b/packages/node/meshnet_node/hardware.py @@ -0,0 +1,33 @@ +"""GPU hardware detection with graceful CPU fallback.""" + +import subprocess + + +def detect_hardware() -> dict: + """Detect GPU model and available VRAM. Returns hardware profile dict.""" + try: + import torch # type: ignore[import] + if torch.cuda.is_available(): + idx = torch.cuda.current_device() + name = torch.cuda.get_device_name(idx) + props = torch.cuda.get_device_properties(idx) + vram_mb = props.total_memory // (1024 * 1024) + return {"device": "cuda", "gpu_name": name, "vram_mb": vram_mb} + except ImportError: + pass + + try: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=5, + ) + if result.returncode == 0 and result.stdout.strip(): + line = result.stdout.strip().splitlines()[0] + parts = line.split(",", 1) + gpu_name = parts[0].strip() + vram_mb = int(parts[1].strip()) if len(parts) > 1 else 0 + return {"device": "cuda", "gpu_name": gpu_name, "vram_mb": vram_mb} + except (FileNotFoundError, subprocess.TimeoutExpired, ValueError, IndexError): + pass + + return {"device": "cpu", "gpu_name": None, "vram_mb": 0} diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py new file mode 100644 index 0000000..c7bc959 --- /dev/null +++ b/packages/node/meshnet_node/startup.py @@ -0,0 +1,154 @@ +"""Full node startup sequence — self-configuring, non-interactive.""" + +from __future__ import annotations + +import json +import socket +import sys +import urllib.error +import urllib.parse +import urllib.request +from pathlib import Path + +from .downloader import download_shard +from .hardware import detect_hardware +from .server import StubNodeServer +from .wallet import load_or_create_wallet + + +def _post_json(url: str, payload: dict, timeout: float = 10.0) -> 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, timeout=timeout) as r: + return json.loads(r.read()) + + +def _get_json(url: str, timeout: float = 10.0) -> dict: + with urllib.request.urlopen(url, timeout=timeout) as r: + return json.loads(r.read()) + + +def run_startup( + tracker_url: str, + port: int = 0, + model: str = "stub-model", + wallet_path: Path | None = None, + cache_dir: Path | None = None, + host: str = "127.0.0.1", + advertise_host: str | None = None, +) -> StubNodeServer: + """Execute the full startup sequence and return a running node server. + + Steps (all non-interactive): + 1. Detect GPU / hardware profile + 2. Load or generate Solana wallet keypair + 3. Query tracker for optimal shard assignment + 4. Download (or stub) the assigned shard + 5. Start local HTTP server + 6. Register with tracker + + Prints a compact status summary on completion. + """ + + tracker_url = tracker_url.rstrip("/") + + # 1. Hardware detection + print("Detecting hardware...", flush=True) + hw = detect_hardware() + device: str = hw["device"] + gpu_name: str | None = hw.get("gpu_name") + vram_mb: int = hw.get("vram_mb", 0) + + if device == "cpu": + print(" WARNING: No CUDA GPU detected — running in CPU mode", flush=True) + else: + print(f" GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True) + + # 2. Wallet + print("Loading wallet...", flush=True) + wallet_kwargs: dict = {} + if wallet_path is not None: + wallet_kwargs["path"] = wallet_path + _, _, address = load_or_create_wallet(**wallet_kwargs) + print(f" Wallet: {address}", flush=True) + + # 3. Shard assignment from tracker + print("Querying tracker for shard assignment...", flush=True) + assign_qs = urllib.parse.urlencode({ + "model": model, + "device": device, + "vram_mb": vram_mb, + }) + try: + assignment = _get_json(f"{tracker_url}/v1/nodes/assign?{assign_qs}") + except urllib.error.URLError as exc: + print(f" ERROR: Cannot reach tracker at {tracker_url}: {exc}", file=sys.stderr, flush=True) + raise + + shard_start: int = assignment["shard_start"] + shard_end: int = assignment["shard_end"] + assigned_model: str = assignment.get("model", model) + hf_repo: str | None = assignment.get("hf_repo") + print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True) + + # 4. Download shard + print("Downloading shard...", flush=True) + dl_kwargs: dict = {} + if cache_dir is not None: + dl_kwargs["cache_dir"] = cache_dir + if hf_repo is not None: + dl_kwargs["hf_repo"] = hf_repo + shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs) + print(f" Cached at: {shard_path}", flush=True) + + # 5. Start HTTP server + is_last = shard_end >= assignment.get("model_layers_end", shard_end) + node = StubNodeServer( + host=host, + port=port, + shard_start=shard_start, + shard_end=shard_end, + is_last_shard=is_last, + ) + actual_port = node.start() + public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) + endpoint = f"http://{public_host}:{actual_port}" + + # 6. Register with tracker + print("Registering with tracker...", flush=True) + try: + reg_resp = _post_json( + f"{tracker_url}/v1/nodes/register", + { + "endpoint": endpoint, + "shard_start": shard_start, + "shard_end": shard_end, + "hardware_profile": hw, + "wallet_address": address, + "score": 1.0, + }, + ) + node_id: str = reg_resp["node_id"] + except Exception: + node.stop() + raise + + # Status summary + hw_str = device.upper() + if gpu_name: + hw_str += f" ({gpu_name}, {vram_mb} MB)" + print( + f"\n{'=' * 32}\n" + f"meshnet-node ready\n" + f" Wallet: {address}\n" + f" Shard: layers {shard_start}-{shard_end} ({assigned_model})\n" + f" Endpoint: {endpoint}\n" + f" Node ID: {node_id}\n" + f" Hardware: {hw_str}\n" + f"{'=' * 32}", + flush=True, + ) + + return node diff --git a/packages/node/meshnet_node/wallet.py b/packages/node/meshnet_node/wallet.py new file mode 100644 index 0000000..3a4477c --- /dev/null +++ b/packages/node/meshnet_node/wallet.py @@ -0,0 +1,64 @@ +"""Solana wallet management — load or generate an Ed25519 keypair. + +Solana keypair format: 64-byte array stored as JSON list of ints. + bytes 0-31: private scalar (seed) + bytes 32-63: public key (on-curve point) +Wallet address: base58-encoded 32-byte public key. +""" + +import json +from pathlib import Path + +_DEFAULT_WALLET_PATH = Path.home() / ".config" / "meshnet" / "wallet.json" + +_B58_ALPHABET = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz" + + +def _b58encode(data: bytes) -> str: + n = int.from_bytes(data, "big") + chars: list[str] = [] + while n: + n, rem = divmod(n, 58) + chars.append(_B58_ALPHABET[rem]) + leading = 0 + for b in data: + if b == 0: + leading += 1 + else: + break + return "1" * leading + "".join(reversed(chars)) + + +def load_or_create_wallet( + path: Path = _DEFAULT_WALLET_PATH, +) -> tuple[bytes, bytes, str]: + """Return (secret_32, public_32, address_base58). + + Loads from *path* if it exists; otherwise generates a new keypair and + saves it. The parent directory is created if needed. + """ + if path.exists(): + mode = path.stat().st_mode & 0o777 + if mode & 0o077: + path.chmod(0o600) + raw = bytes(json.loads(path.read_text())) + if len(raw) != 64: + raise ValueError(f"Wallet at {path} has unexpected length {len(raw)}, expected 64") + secret, public = raw[:32], raw[32:] + return secret, public, _b58encode(public) + + from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey + from cryptography.hazmat.primitives.serialization import ( + Encoding, NoEncryption, PrivateFormat, PublicFormat, + ) + + priv = Ed25519PrivateKey.generate() + secret = priv.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption()) + public = priv.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw) + keypair = list(secret) + list(public) + + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(keypair)) + path.chmod(0o600) + + return secret, public, _b58encode(public) diff --git a/packages/node/pyproject.toml b/packages/node/pyproject.toml index 68b4cf3..46225c8 100644 --- a/packages/node/pyproject.toml +++ b/packages/node/pyproject.toml @@ -8,6 +8,11 @@ version = "0.1.0" description = "Distributed Inference Network node client" requires-python = ">=3.10" +dependencies = [ + "cryptography>=41", + "huggingface-hub>=0.20", +] + [project.scripts] meshnet-node = "meshnet_node.cli:main" diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 785e3da..dd66e8e 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -4,7 +4,8 @@ HTTP API contract: - POST /v1/nodes/register Request JSON: { "endpoint": "http://host:port", "shard_start": int, "shard_end": int, - "hardware_profile": object, "score": number optional + "hardware_profile": object, "wallet_address": str optional, + "score": number optional } Response 200: {"node_id": str} Response 400: {"error": str} @@ -32,7 +33,7 @@ DEFAULT_MODEL_PRESETS: dict[str, dict] = { class _NodeEntry: __slots__ = ( "node_id", "endpoint", "shard_start", "shard_end", - "hardware_profile", "score", "last_heartbeat", + "hardware_profile", "wallet_address", "score", "last_heartbeat", ) def __init__( @@ -42,6 +43,7 @@ class _NodeEntry: shard_start: int, shard_end: int, hardware_profile: dict, + wallet_address: str | None, score: float, ) -> None: self.node_id = node_id @@ -49,6 +51,7 @@ class _NodeEntry: self.shard_start = shard_start self.shard_end = shard_end self.hardware_profile = hardware_profile + self.wallet_address = wallet_address self.score = score self.last_heartbeat: float = time.monotonic() @@ -146,6 +149,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): parsed = urllib.parse.urlparse(self.path) if parsed.path == "/v1/route": self._handle_route(parsed) + elif parsed.path == "/v1/nodes/assign": + self._handle_assign(parsed) else: self.send_response(404) self.end_headers() @@ -180,6 +185,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if not isinstance(hardware_profile, dict): self._send_json(400, {"error": "hardware_profile must be an object"}) return + wallet_address = body.get("wallet_address") + if wallet_address is not None and not isinstance(wallet_address, str): + self._send_json(400, {"error": "wallet_address must be a string"}) + return node_id = str(uuid.uuid4()) entry = _NodeEntry( @@ -188,6 +197,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): shard_start=shard_start, shard_end=shard_end, hardware_profile=hardware_profile, + wallet_address=wallet_address, score=score, ) with server.lock: @@ -207,6 +217,81 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): entry.last_heartbeat = time.monotonic() self._send_json(200, {}) + def _handle_assign(self, parsed: urllib.parse.ParseResult): + """Return an optimal shard assignment for a node given its hardware profile. + + Query params: + model — model preset name (default: first preset) + device — "cuda" | "cpu" + vram_mb — integer VRAM in MB (0 for CPU) + + The greedy strategy: find the first gap in current layer coverage + and assign it. If no gap exists, assign the full model range so the + node provides redundant coverage. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + params = urllib.parse.parse_qs(parsed.query) + + model_list = params.get("model") + if not model_list: + model = next(iter(server.model_presets), None) + if model is None: + self._send_json(503, {"error": "no model presets configured"}) + return + else: + model = model_list[0] + + preset = server.model_presets.get(model) + if preset is None: + self._send_json(404, {"error": f"unknown model preset: {model!r}"}) + return + + required_start: int = preset["layers_start"] + required_end: int = preset["layers_end"] + + with server.lock: + self._purge_expired_nodes() + alive = list(server.registry.values()) + + device = params.get("device", ["cpu"])[0] + try: + vram_mb = int(params.get("vram_mb", ["0"])[0]) + except ValueError: + vram_mb = 0 + max_layers = required_end - required_start + 1 + if device != "cuda" or vram_mb < 8192: + max_layers = min(max_layers, 16) + + # Collect covered intervals sorted by start layer. + covered = sorted( + [(n.shard_start, n.shard_end) for n in alive], + key=lambda t: t[0], + ) + + # Walk from required_start to find the first uncovered layer. + gap_start = required_start + for s, e in covered: + if s <= gap_start: + gap_start = max(gap_start, e + 1) + else: + break # gap found before this node + + if gap_start > required_end: + # Full coverage exists — assign the full range for redundancy. + shard_start = required_start + shard_end = min(required_end, shard_start + max_layers - 1) + else: + shard_start = gap_start + shard_end = min(required_end, shard_start + max_layers - 1) + + self._send_json(200, { + "shard_start": shard_start, + "shard_end": shard_end, + "model": model, + "model_layers_end": required_end, + **({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}), + }) + def _handle_route(self, parsed: urllib.parse.ParseResult): server: _TrackerHTTPServer = self.server # type: ignore[assignment] params = urllib.parse.parse_qs(parsed.query) diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py new file mode 100644 index 0000000..9eefc73 --- /dev/null +++ b/tests/test_node_startup.py @@ -0,0 +1,297 @@ +"""US-004 integration tests: node self-configuring startup sequence.""" + +import json +import sys +import types +import urllib.request +from pathlib import Path + +import pytest + +from meshnet_node.downloader import download_shard +from meshnet_node.hardware import detect_hardware +from meshnet_node.startup import run_startup +from meshnet_node.wallet import _b58encode, load_or_create_wallet +from meshnet_tracker.server import TrackerServer + + +# --------------------------------------------------------------------------- +# Unit tests — hardware, wallet, downloader +# --------------------------------------------------------------------------- + + +def test_detect_hardware_returns_valid_profile(): + """Hardware detection always returns a dict with required keys.""" + hw = detect_hardware() + assert hw["device"] in {"cuda", "cpu"} + assert isinstance(hw.get("vram_mb"), int) + if hw["device"] == "cpu": + assert hw["gpu_name"] is None + assert hw["vram_mb"] == 0 + else: + assert isinstance(hw["gpu_name"], str) and hw["gpu_name"] + assert hw["vram_mb"] > 0 + + +def test_wallet_generates_new_keypair(tmp_path): + """A new wallet is created when none exists, saved to disk.""" + wallet_file = tmp_path / "wallet.json" + assert not wallet_file.exists() + + secret, public, address = load_or_create_wallet(path=wallet_file) + + assert wallet_file.exists() + assert wallet_file.stat().st_mode & 0o777 == 0o600 + assert len(secret) == 32 + assert len(public) == 32 + assert len(address) > 20 # base58-encoded 32 bytes is 43-44 chars typically + # Solana addresses are base58 — no '0', 'O', 'I', 'l' + for ch in "0OIl": + assert ch not in address, f"Invalid base58 char {ch!r} in address {address!r}" + + +def test_wallet_loads_existing_keypair(tmp_path): + """Loading the same wallet file twice returns identical keys and address.""" + wallet_file = tmp_path / "wallet.json" + + secret1, public1, address1 = load_or_create_wallet(path=wallet_file) + secret2, public2, address2 = load_or_create_wallet(path=wallet_file) + + assert secret1 == secret2 + assert public1 == public2 + assert address1 == address2 + + +def test_wallet_load_repairs_insecure_permissions(tmp_path): + """Existing private key files are tightened to owner-only permissions.""" + wallet_file = tmp_path / "wallet.json" + load_or_create_wallet(path=wallet_file) + wallet_file.chmod(0o644) + + load_or_create_wallet(path=wallet_file) + + assert wallet_file.stat().st_mode & 0o777 == 0o600 + + +def test_base58_counts_only_leading_zero_bytes(): + """Zero bytes inside the public key do not become extra base58 leading ones.""" + assert _b58encode(bytes([0, 1, 0])) == "15R" + + +def test_download_shard_stub_creates_cache(tmp_path): + """Stub-model shard creates a local cache file without network access.""" + shard_dir = download_shard("stub-model", 0, 31, cache_dir=tmp_path) + assert shard_dir.exists() + weights = shard_dir / "weights.json" + assert weights.exists() + data = json.loads(weights.read_text()) + assert data["stub"] is True + assert data["shard_start"] == 0 + assert data["shard_end"] == 31 + + +def test_download_shard_uses_huggingface_when_repo_is_assigned(tmp_path, monkeypatch): + """Non-stub shards use the HuggingFace snapshot_download path.""" + calls = [] + + def fake_snapshot_download(repo_id, cache_dir, local_dir): + calls.append({"repo_id": repo_id, "cache_dir": cache_dir, "local_dir": local_dir}) + Path(local_dir).mkdir(parents=True, exist_ok=True) + return local_dir + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=fake_snapshot_download), + ) + + shard_dir = download_shard( + "tiny-llama", + 0, + 3, + cache_dir=tmp_path, + hf_repo="org/tiny-llama-shards", + progress=False, + ) + + assert shard_dir == tmp_path / "tiny-llama" / "layers_0-3" + assert calls == [{ + "repo_id": "org/tiny-llama-shards", + "cache_dir": str(tmp_path), + "local_dir": str(shard_dir), + }] + + +def test_download_shard_stub_idempotent(tmp_path): + """Calling download_shard twice does not error — file already exists.""" + download_shard("stub-model", 0, 31, cache_dir=tmp_path, progress=False) + shard_dir = download_shard("stub-model", 0, 31, cache_dir=tmp_path, progress=False) + assert shard_dir.exists() + + +# --------------------------------------------------------------------------- +# Tracker assign endpoint +# --------------------------------------------------------------------------- + + +def _get_json(url: str) -> dict: + with urllib.request.urlopen(url, timeout=5) as r: + return json.loads(r.read()) + + +def test_tracker_assign_returns_shard_for_empty_registry(): + """Tracker assigns the full layer range when no nodes are registered.""" + tracker = TrackerServer() + port = tracker.start() + try: + resp = _get_json( + f"http://127.0.0.1:{port}/v1/nodes/assign?model=stub-model&device=cpu&vram_mb=0" + ) + assert resp["shard_start"] == 0 + assert resp["shard_end"] == 15 + assert resp["model"] == "stub-model" + assert resp["model_layers_end"] == 31 + finally: + tracker.stop() + + +def test_tracker_assign_fills_gap(): + """Tracker assigns the first uncovered layer range when a node is already registered.""" + import json as _json + import urllib.request as _ur + + tracker = TrackerServer() + port = tracker.start() + try: + # Register a node covering layers 0-15 + data = _json.dumps({ + "endpoint": "http://127.0.0.1:9100", + "shard_start": 0, + "shard_end": 15, + "hardware_profile": {}, + "score": 1.0, + }).encode() + req = _ur.Request( + f"http://127.0.0.1:{port}/v1/nodes/register", + data=data, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with _ur.urlopen(req) as r: + r.read() + + # Assignment should fill the gap: layers 16-31 + resp = _get_json( + f"http://127.0.0.1:{port}/v1/nodes/assign?model=stub-model&device=cpu&vram_mb=0" + ) + assert resp["shard_start"] == 16 + assert resp["shard_end"] == 31 + finally: + tracker.stop() + + +def test_tracker_assign_returns_huggingface_repo_when_configured(): + """Tracker includes the HuggingFace repo identifier in shard assignments.""" + tracker = TrackerServer(model_presets={ + "tiny-llama": {"layers_start": 0, "layers_end": 7, "hf_repo": "org/tiny-llama-shards"} + }) + port = tracker.start() + try: + resp = _get_json( + f"http://127.0.0.1:{port}/v1/nodes/assign?model=tiny-llama&device=cuda&vram_mb=24576" + ) + assert resp["model"] == "tiny-llama" + assert resp["hf_repo"] == "org/tiny-llama-shards" + assert resp["shard_start"] == 0 + assert resp["shard_end"] == 7 + finally: + tracker.stop() + + +# --------------------------------------------------------------------------- +# Full startup integration test +# --------------------------------------------------------------------------- + + +def test_full_startup_sequence(tmp_path): + """Full startup: hardware → wallet → assign → download → start → register.""" + tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}}) + tracker_port = tracker.start() + tracker_url = f"http://127.0.0.1:{tracker_port}" + + wallet_path = tmp_path / "wallet.json" + cache_dir = tmp_path / "shards" + + try: + node = run_startup( + tracker_url=tracker_url, + model="stub-model", + wallet_path=wallet_path, + cache_dir=cache_dir, + ) + try: + # Wallet was created on disk + assert wallet_path.exists() + wallet_data = json.loads(wallet_path.read_text()) + assert len(wallet_data) == 64 + + # Shard was cached + assert any(cache_dir.rglob("weights.json")) + + # Node appears in tracker registry (route resolves successfully) + route_resp = _get_json(f"{tracker_url}/v1/route?model=stub-model") + assert len(route_resp["route"]) >= 1 + assert node.port is not None + assert any(str(node.port) in ep for ep in route_resp["route"]) + + # Node accepts an inference request + payload = json.dumps({ + "messages": [{"role": "user", "content": "hello"}], + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{node.port}/v1/infer", + data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5) as resp: + assert resp.status == 200 + body = json.loads(resp.read()) + # Last-shard node returns text; first-shard returns activations + assert "text" in body or "activations" in body + finally: + node.stop() + finally: + tracker.stop() + + +def test_startup_cpu_fallback(tmp_path, monkeypatch): + """Node starts with CPU warning when no GPU is detected.""" + import meshnet_node.startup as startup_mod + + monkeypatch.setattr( + startup_mod, + "detect_hardware", + lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0}, + ) + + tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}}) + tracker_port = tracker.start() + tracker_url = f"http://127.0.0.1:{tracker_port}" + + try: + node = run_startup( + tracker_url=tracker_url, + model="stub-model", + wallet_path=tmp_path / "wallet.json", + cache_dir=tmp_path / "shards", + ) + try: + # Node is running even on CPU + assert node.port is not None + route_resp = _get_json(f"{tracker_url}/v1/route?model=stub-model") + assert len(route_resp["route"]) >= 1 + finally: + node.stop() + finally: + tracker.stop()