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

inference working

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-29 23:54:35 +03:00
parent 607d49f5b0
commit 1bdfce657d
16 changed files with 1899 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.claude/worktrees/feat+us-016 Submodule

Submodule .claude/worktrees/feat+us-016 added at 080d49b2c2

206
.scratch/distributed-inference-network/issues/16-mining-cli-ux.md Normal file
View File

@@ -0,0 +1,206 @@
# US-016 — Mining-style node startup CLI + live dashboard
## Goal
Replace the bare flag-driven `meshnet-node start` with a wizard-guided first-run experience modelled on GPU mining clients (like PhoenixMiner, lolMiner, etc.). After the wizard, the terminal switches to a live status dashboard showing real-time node health and earnings.
## Wizard flow (first run only)
```
╔══════════════════════════════════════════════════════════╗
║ meshnet-node v0.1.0 ║
║ Distributed AI Inference — Node Setup ║
╚══════════════════════════════════════════════════════════╝
Detecting hardware...
GPU 0: NVIDIA RTX 4090 24 GB VRAM ✓
GPU 1: NVIDIA RTX 3090 24 GB VRAM ✓
Select a model to serve:
# Model Layers NF4 INT8 BF16
1 Llama-3-70B-Instruct 80 ✓18GB ✓40GB ✗80GB
2 Qwen-2.5-72B-Instruct 80 ✓19GB ✗41GB ✗81GB
3 Mixtral-8x7B-Instruct-v0.1 32 ✓ 7GB ✓14GB ✓27GB
4 Phi-3-medium-128k-instruct 40 ✓ 4GB ✓ 8GB ✓15GB
5 [Browse HuggingFace…]
Enter number [1]: _
Quantization [nf4/int8/bf16] (nf4 recommended for 24GB): _
Download directory [~/.meshnet/models]: _
Tracker URL [http://localhost:8080]: _
Wallet path [~/.config/meshnet/wallet.json] (new wallet will be created): _
Config saved to ~/.config/meshnet/config.json
Starting node…
```
Second run with existing config:
```
meshnet-node
Reading config from ~/.config/meshnet/config.json
Model: Llama-3-70B-Instruct Quant: nf4 Shard: layers 0–15
Tracker: http://192.168.1.10:8080
Starting…
```
## Live dashboard (once running)
Renders every 2 seconds using `rich.live`. Fallback: plain-text status line if `rich` is unavailable or terminal is not a TTY (important for WSL2 / SSH).
```
meshnet-node Llama-3-70B-Instruct [nf4] shard 0–15/80 up 00:03:22
GPU 0 RTX 4090 GPU ████████░░ 73% VRAM 18.2/24.0 GB 45°C
GPU 1 RTX 3090 GPU ███░░░░░░░ 28% VRAM 8.7/24.0 GB 38°C
Tokens/sec ▁▂▃▄▅▆▇█ 42.3 t/s (EMA 30s)
Requests 1,247 served 3 active
Peers 8 connected (tracker: ✓ relay: ✓)
TAI earned 0.00 TAI (payments active after US-006)
Uptime 00:03:22
[q] quit [r] reset stats [c] compact view
```
Compact mode (`--compact` or pressing `c`) shows a single status line:
```
[43t/s VRAM18.2GB req1247 peers8 up3m22s]
```
## Implementation notes
### Hardware detection
```python
import torch
def detect_gpus() -> list[dict]:
gpus = []
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
props = torch.cuda.get_device_properties(i)
gpus.append({
"index": i,
"name": props.name,
"vram_gb": props.total_memory / 1e9,
"backend": "cuda"
})
# ROCm / Apple Silicon stubs for later
return gpus
```
### Curated model list
`packages/node/meshnet_node/model_catalog.py` — a hardcoded list of `ModelPreset` dataclasses:
```python
@dataclass
class ModelPreset:
name: str # display name
hf_repo: str # HuggingFace repo ID
num_layers: int
vram_gb: dict # {"nf4": 18, "int8": 40, "bf16": 80}
description: str # one-line description
```
Initial list (expand over time):
- `meta-llama/Meta-Llama-3-70B-Instruct` — 80L, NF4 18GB, INT8 40GB, BF16 80GB
- `Qwen/Qwen2.5-72B-Instruct` — 80L, NF4 19GB, INT8 41GB, BF16 81GB
- `mistralai/Mixtral-8x7B-Instruct-v0.1` — 32L, NF4 7GB, INT8 14GB, BF16 27GB
- `microsoft/Phi-3-medium-128k-instruct` — 40L, NF4 4GB, INT8 8GB, BF16 15GB
- `google/gemma-2-27b-it` — 46L, NF4 10GB, INT8 20GB, BF16 40GB
### HuggingFace Browse
```python
from huggingface_hub import list_models
def browse_hf(top_n=20) -> list[dict]:
models = list_models(
pipeline_tag="text-generation",
library="transformers",
sort="downloads",
direction=-1,
limit=top_n,
cardData=True,
)
return [{"repo": m.modelId, "downloads": m.downloads} for m in models]
```
### Persistent config
`~/.config/meshnet/config.json`:
```json
{
"model_hf_repo": "meta-llama/Meta-Llama-3-70B-Instruct",
"quantization": "nf4",
"download_dir": "~/.meshnet/models",
"tracker_url": "http://192.168.1.10:8080",
"wallet_path": "~/.config/meshnet/wallet.json",
"shard_start": null,
"shard_end": null,
"updatedAt": "2026-06-29T..."
}
```
`shard_start`/`shard_end`: null means tracker auto-assigns. User can pin a range for dedicated partial-model nodes.
### CLI flags
All wizard answers are overridable without re-running the wizard:
```
meshnet-node [start]
--model <hf-repo-id> # e.g. meta-llama/Meta-Llama-3-70B-Instruct
--quantization [bf16|int8|nf4]
--download-dir <path>
--tracker <url>
--wallet <path>
--shard-start <int> # pin shard range (optional)
--shard-end <int>
--reset-config # ignore saved config, re-run wizard
--no-tui # plain-text output (for CI / headless)
--compact # single-line status instead of full dashboard
meshnet-node models # list curated models and exit
meshnet-node models --browse # list HF Hub top-20 and exit
meshnet-node config # print current config and exit
```
### WSL2 / non-TTY fallback
```python
import sys, os
def is_interactive_tty() -> bool:
return sys.stdout.isatty() and os.environ.get("TERM") not in ("dumb", "")
if not is_interactive_tty():
# fall back to plain-text periodic status
run_plain_status_loop(node)
else:
run_rich_dashboard(node)
```
Do NOT use `termios`, `fcntl`, or `/dev/tty` — these break in Windows cmd.exe and some WSL2 terminal emulators.
## Acceptance criteria
- `meshnet-node` with no args and no config → wizard starts
- Wizard detects GPU and marks `[too large]` for models that exceed available VRAM
- `meshnet-node models` prints curated list and exits
- `meshnet-node models --browse` calls HF Hub API, prints top-20, exits
- Second run (config exists) → skips wizard, starts immediately
- `--reset-config` re-runs wizard even with config present
- All wizard inputs override-able via CLI flags
- Live rich dashboard renders and updates every 2s when running in a TTY
- Falls back to plain-text when not a TTY (CI / WSL2 without TERM set)
- Ctrl-C prints a clean summary line and exits 0
- `python -m pytest` passes from repo root
- Commit only this story's changes

205
.scratch/distributed-inference-network/issues/17-p2p-gossip-relay-ssl.md Normal file
View File

@@ -0,0 +1,205 @@
# US-017 — P2P gossip, NAT-traversal relay node, and SSL/TLS
## Goal
Nodes must work behind NAT (home routers, cloud VMs without public IPs) and must communicate securely. Implement:
1. **SSL/TLS everywhere** — all HTTP between nodes/tracker is HTTPS; all WebSocket gossip is WSS
2. **mDNS peer discovery** — nodes on the same LAN find each other automatically (no config)
3. **WebSocket gossip PubSub** — nodes propagate join/leave/coverage-update events in near-real-time
4. **Circuit relay node** — team-run public relay (`packages/relay`) that enables NAT traversal and bootstraps new nodes joining from the internet
Architecture is designed to migrate to libp2p GossipSub + Kademlia DHT without breaking the message schema (topic names and payload formats are stable contracts).
## Gossip protocol
### Transport
WebSocket (`wss://`) using the `websockets` Python library. Each node maintains persistent WSS connections to:
- The relay node (always, bootstraps peer list)
- Up to 8 direct peers (Kademlia-style target fanout; peers discovered via mDNS + relay peer list)
### Topics
All messages are JSON with an envelope:
```json
{
"topic": "node-join",
"version": 1,
"from_peer": "<peer_id>",
"timestamp": "<iso8601>",
"payload": { ... }
}
```
| Topic | Direction | Payload |
|-------|-----------|---------|
| `node-join` | broadcast | `{peer_id, addr, models: [{model_preset, shard_start, shard_end}], vram_gb, quant}` |
| `node-leave` | broadcast | `{peer_id, reason}` |
| `coverage-update` | broadcast | `{model_preset, coverage: [{start, end, count}]}` |
| `heartbeat` | peer→relay | `{peer_id, addr, uptime_s, tokens_per_sec}` |
| `peer-list` | relay→peer | `{peers: [{peer_id, addr}]}` |
| `relay-announce` | relay→all | `{relay_id, relay_url, capacity}` |
Gossip fanout: each node re-broadcasts received messages to all its peers (simple flooding with `seen_ids` dedup, TTL=3 hops). Migration to GossipSub mesh routing is a later ADR.
### Peer ID
`peer_id = sha256(public_key)[:16].hex()` — generated on first run, stored in `~/.config/meshnet/identity.json`. The same keypair is used for TLS client certificates (mTLS) in future work.
## mDNS LAN discovery
Use Python `zeroconf` library. Service type: `_meshnet._tcp.local.`
```python
from zeroconf import ServiceInfo, Zeroconf
info = ServiceInfo(
"_meshnet._tcp.local.",
f"{peer_id}._meshnet._tcp.local.",
addresses=[socket.inet_aton(local_ip)],
port=node_port,
properties={"peer_id": peer_id, "version": "1"},
)
zc = Zeroconf()
zc.register_service(info)
```
On startup, nodes also browse for `_meshnet._tcp.local.` to discover existing nodes. mDNS is LAN-only (does not traverse routers), which is correct for LAN discovery.
## NAT traversal: circuit relay
### How it works
1. Node A (behind NAT) cannot accept inbound TCP connections
2. Node A connects outbound to the public relay via WSS
3. Node A tells the tracker: `"effective_addr": "wss://relay.meshnet.ai/relay/{peer_id_A}"`
4. Node B (wants to call A) connects to the relay at the above URL
5. Relay proxies the TCP stream between A and B
Hole-punching (direct connection via STUN) is attempted first (future work). Relay is the fallback.
### meshnet-relay
`packages/relay/meshnet_relay/server.py` — a standalone aiohttp server:
```
GET /health → {status: ok}
GET /v1/peers → [{peer_id, addr, last_seen}]
POST /v1/gossip → receive a gossip message, fan out to connected peers
WSS /ws → persistent gossip connection (subscribe to all topics)
WSS /relay/{peer_id} → circuit relay proxy to that peer_id
GET /v1/relay/capacity → {connected_peers: N, max_peers: 500}
```
CLI:
```
meshnet-relay [--port 8443] [--cert path/to/cert.pem] [--key path/to/key.pem]
[--tracker-url http://...] [--max-peers 500]
```
The relay can optionally proxy to the tracker (so `relay.meshnet.ai` is the single internet-visible endpoint).
## SSL/TLS setup
### Node certificate (self-signed, auto-generated)
On first run, `meshnet-node` generates a self-signed RSA-2048 cert valid for 10 years:
```python
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
```
Cert saved to `~/.config/meshnet/node_cert.pem` + `node_key.pem`. Fingerprint stored in config and shared with tracker via heartbeat. Nodes connecting to each other validate the fingerprint (TOFU — trust on first use), not the CA chain.
### Relay certificate
The relay uses a real Let's Encrypt cert (cert-bot or acme.sh). The relay cert is pinned in `packages/p2p/relay_bootstrap.json`:
```json
{
"relays": [
{
"url": "wss://relay.meshnet.ai:8443",
"cert_fingerprint": "sha256:<hex>",
"operator": "meshnet-team"
}
]
}
```
### All HTTP switched to HTTPS
`meshnet-node` starts an HTTPS server using `ssl.SSLContext`. `meshnet-tracker` similarly. All outbound `httpx` / `aiohttp` calls use TLS verification against pinned fingerprints (not the system CA store — too many corporate proxies break this).
## Tracker changes
Heartbeat payload gains new fields:
```json
{
"peer_id": "a1b2c3d4e5f6a1b2",
"effective_addr": "https://192.168.1.42:8001",
"relay_addr": "wss://relay.meshnet.ai:8443/relay/a1b2c3d4e5f6a1b2",
"cert_fingerprint": "sha256:...",
"gossip_peers": ["peer_id_1", "peer_id_2"]
}
```
Tracker uses `effective_addr` (direct) or `relay_addr` (fallback) when building inference routes.
## Integration test
```
tests/test_gossip_and_relay.py
scenario:
1. Start a local relay (localhost:18443)
2. Start node A (no inbound port — simulate NAT by binding to 127.0.0.1 only)
3. Start node B (public-reachable on localhost)
4. Both register with relay; relay peer-list includes both
5. Node B sends a gossip node-join message
6. Assert node A receives it within 500ms
7. Start tracker; confirm tracker's node registry includes node A via relay_addr
8. Send inference request; assert it routes through relay to node A
```
## Package layout
```
packages/relay/
pyproject.toml
meshnet_relay/
__init__.py
server.py # aiohttp relay + gossip hub + circuit relay proxy
cli.py # meshnet-relay entrypoint
peer_registry.py # in-memory {peer_id: {addr, last_seen, ...}}
circuit_relay.py # WSS proxy between two peers
packages/p2p/
meshnet_p2p/
gossip.py # GossipClient — connect to relay + peers, pub/sub
mdns.py # ZeroconfDiscovery — mDNS announce + browse
identity.py # PeerIdentity — generate/load peer_id + keypair
tls.py # cert generation, fingerprint, SSLContext helpers
packages/node/meshnet_node/
gossip_integration.py # wires GossipClient into node lifecycle
```
## Acceptance criteria
- All node↔node and node↔tracker HTTP uses HTTPS; self-signed cert auto-generated on first run
- `cert_fingerprint` included in heartbeat; tracker stores and logs it
- mDNS: two nodes on the same LAN discover each other without manual tracker URL (test with two localhost processes using different mDNS names)
- Relay: `meshnet-relay` starts, accepts WSS connections, fans out gossip messages to all connected peers
- Circuit relay: node A (127.0.0.1-only) can receive a gossip message via the relay from node B
- Tracker routes inference to node A using `relay_addr` when direct addr not reachable
- `relay_bootstrap.json` exists in `packages/p2p/` with at least one entry (localhost for tests)
- ADR-0010 documents the gossip architecture and libp2p migration path
- `python -m pytest` passes from repo root
- Commit only this story's changes

157
.scratch/distributed-inference-network/issues/18-two-machine-lan-test.md Normal file
View File

@@ -0,0 +1,157 @@
# US-018 — End-to-end two-machine LAN inference test
## Goal
Run real distributed inference across two physical machines: the Linux rig and a Windows 11 rig running WSL2. Document every setup step, firewall rule, and gotcha so this is repeatable. The test script exits 0 with token output and timing, proving the network works.
## Network topology for LAN test
```
[Linux machine] [Windows 11 / WSL2]
meshnet-tracker :8080 meshnet-node (shard B)
meshnet-node :8001 (shard A, tracker-mode)
meshnet-gateway :8000 (optional, for OpenAI-compat)
Client (either machine):
scripts/test_lan_inference.py --tracker http://192.168.1.10:8080
```
The Linux machine runs the tracker + the first-shard node (tracker-mode). The Windows/WSL2 machine runs the second-shard node. A small model (e.g. Phi-3-medium at BF16, fits on one GPU each) is split across both.
## WSL2 setup (Windows side)
`docs/INSTALL_WINDOWS.md` covers:
1. Enable WSL2: `wsl --install -d Ubuntu-24.04`
2. CUDA in WSL2: install NVIDIA driver on Windows (NOT inside WSL); WSL2 gets CUDA automatically
- Verify: `nvidia-smi` inside WSL2 should show GPU
3. Install Python 3.11+ and pip inside WSL2
4. `pip install -e packages/node packages/p2p` (clone repo first)
5. Firewall: Windows Defender must allow inbound WSL2 → LAN on node port
- PowerShell: `New-NetFirewallRule -DisplayName "meshnet-node" -Direction Inbound -Protocol TCP -LocalPort 8001 -Action Allow`
6. WSL2 IP: WSL2 has its own NAT'd IP (172.x.x.x); to expose to LAN, either:
- Option A: `netsh interface portproxy add v4tov4 listenport=8001 listenaddress=0.0.0.0 connectport=8001 connectaddress=$(wsl hostname -I)`
- Option B: use the relay node (US-017) — no port forwarding needed
## Linux setup
Standard install (already done after US-016). Firewall:
```bash
# If using ufw
sudo ufw allow 8080/tcp # tracker
sudo ufw allow 8001/tcp # node
sudo ufw allow 8000/tcp # gateway (optional)
```
## Model split
For the test, use a model that has enough layers to split meaningfully but fits comfortably in memory. Phi-3-medium-128k-instruct (40 layers, BF16 15GB) works on a single 24GB GPU on each machine:
- Linux node: layers 0–19 (tracker-mode, owns tokenizer + embed_tokens)
- Windows/WSL2 node: layers 20–39
Start sequence:
```bash
# Terminal 1 (Linux) — tracker
meshnet-tracker --port 8080
# Terminal 2 (Linux) — first-shard node (tracker-mode auto-detected because shard_start=0)
meshnet-node --model microsoft/Phi-3-medium-128k-instruct \
--quantization bf16 \
--shard-start 0 --shard-end 19 \
--tracker http://localhost:8080 \
--port 8001
# Terminal 3 (Windows WSL2) — second-shard node
meshnet-node --model microsoft/Phi-3-medium-128k-instruct \
--quantization bf16 \
--shard-start 20 --shard-end 39 \
--tracker http://192.168.1.10:8080 \
--port 8001
```
## Test script
`scripts/test_lan_inference.py`:
```python
#!/usr/bin/env python3
"""
End-to-end LAN inference test.
Usage: python scripts/test_lan_inference.py --tracker http://192.168.1.10:8080
"""
import argparse, time, httpx, json
MESSAGES = [
{"role": "user", "content": "What is 7 × 8? Answer in one word."},
{"role": "user", "content": "Name the capital of France in one word."},
{"role": "user", "content": "Complete the sequence: 1, 1, 2, 3, 5, ___"},
]
def run_test(tracker_url: str, gateway_url: str | None):
# Discover inference entry point via tracker if gateway not given
if not gateway_url:
r = httpx.get(f"{tracker_url}/v1/tracker-nodes/phi-3-medium", timeout=5)
r.raise_for_status()
nodes = r.json()
assert nodes, "No tracker-mode nodes registered — is the first-shard node running?"
gateway_url = nodes[0]["url"]
print(f"Inference endpoint: {gateway_url}")
print(f"Tracker: {tracker_url}")
print()
for i, msg in enumerate(MESSAGES):
t0 = time.monotonic()
r = httpx.post(
f"{gateway_url}/v1/chat/completions",
json={"model": "phi-3-medium", "messages": [msg], "stream": False},
timeout=60,
)
r.raise_for_status()
data = r.json()
elapsed = time.monotonic() - t0
content = data["choices"][0]["message"]["content"]
tokens = data["usage"]["completion_tokens"]
tps = tokens / elapsed if elapsed > 0 else 0
print(f"[{i+1}] Q: {msg['content']}")
print(f" A: {content}")
print(f" {tokens} tokens {elapsed:.2f}s {tps:.1f} t/s")
print()
print("✓ All 3 requests completed successfully")
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--tracker", required=True)
p.add_argument("--gateway", default=None)
args = p.parse_args()
run_test(args.tracker, args.gateway)
```
## Docs: TWO_MACHINE_TEST.md
`docs/TWO_MACHINE_TEST.md` must cover:
1. Prerequisites (models downloaded on both machines, same model ID, complementary shard ranges)
2. Start order: tracker first, then nodes, then test script
3. How to verify nodes are registered: `GET /v1/nodes` on tracker
4. How to verify coverage: `GET /v1/coverage/phi-3-medium` — all 40 layers must show node_count ≥ 1
5. How to run the test script
6. Expected output
7. Latency breakdown: how to read per-hop latency from node logs
8. **Known Issues** section — updated during actual test run with real gotchas
## Acceptance criteria
- `docs/INSTALL_WINDOWS.md` covers WSL2 + CUDA + meshnet-node install end-to-end
- `docs/TWO_MACHINE_TEST.md` covers the full two-machine setup and test procedure
- `scripts/test_lan_inference.py` exists and is executable
- When run against a real two-machine LAN setup: script exits 0, prints 3 valid answers with timing
- Coverage map shows 100% coverage (no gap) after both nodes register
- Known Issues section in TWO_MACHINE_TEST.md contains at least the issues encountered during this test run
- No new pytest failures from repo root (this story adds docs + a script, not new Python packages)
- Commit only this story's changes

67
docs/adr/0010-p2p-gossip-and-nat-relay.md Normal file
View File

@@ -0,0 +1,67 @@
# ADR-0010: P2P gossip, NAT-traversal relay, and TLS
## Status: Accepted
## Context
All node-to-node and node-to-tracker communication in the prototype is plain HTTP over a LAN or direct-IP internet connection. This has three problems:
1. **NAT blocking**: Most home and cloud nodes cannot accept inbound TCP connections.
2. **No encryption**: Activations and heartbeats are in plaintext.
3. **Polling overhead**: Nodes poll the tracker for coverage changes every 30s. This is slow to react to node churn and does not scale past a few hundred nodes.
The reference implementation (Petals) solves this with libp2p — GossipSub for pub/sub and Kademlia DHT for peer discovery. We adopt the same goals but start with simpler, more stable building blocks that can be swapped for libp2p later without changing the message schema.
## Decisions
### 1. TLS everywhere
All HTTP between nodes, tracker, and gateway uses HTTPS (TLS 1.3). Self-signed certificates are auto-generated on first node start and stored in `~/.config/meshnet/`. The certificate fingerprint is included in every heartbeat and gossip envelope. Nodes use TOFU (trust on first use) — they accept a peer's cert on first contact and pin the fingerprint; connections from the same peer with a different fingerprint are rejected.
The relay node uses a real CA-signed certificate (Let's Encrypt) because it is the internet-facing bootstrap point.
### 2. mDNS for LAN peer discovery
Python `zeroconf` library. Service type: `_meshnet._tcp.local.`. A node announces itself on startup and browses for existing peers. This is zero-config discovery for home and lab networks. mDNS does not traverse routers, which is correct — LAN discovery should not bleed into the internet.
### 3. WebSocket PubSub for gossip
Each node maintains persistent WSS connections to the relay and up to 8 direct peers. Messages use a stable JSON envelope with a `topic`, `version`, `from_peer`, and `payload`. Topics: `node-join`, `node-leave`, `coverage-update`, `heartbeat`, `peer-list`, `relay-announce`.
Simple flooding with `seen_ids` dedup and TTL=3 is good enough for the prototype. The message schema is stable; the fanout mechanism can be replaced with GossipSub mesh routing without changing the schema.
### 4. Circuit relay node for NAT traversal
A team-operated public relay (`packages/relay`, CLI: `meshnet-relay`) is the internet bootstrap point. A node behind NAT:
1. Connects outbound to the relay via WSS
2. Advertises `relay_addr = wss://relay.meshnet.ai:8443/relay/{peer_id}` to the tracker
3. Other nodes proxy connections through the relay when the direct addr is not reachable
Hole-punching (STUN + simultaneous TCP open) is deferred to a future story. Circuit relay is the reliable fallback.
The relay is stateless in terms of inference — it only proxies bytes. It does not decrypt activations.
### 5. Bootstrap peer list
`packages/p2p/relay_bootstrap.json` contains the team-operated relay endpoints with their TLS fingerprints. New nodes load this file on startup to find their first peer. The list is bundled with the package and updated via pip upgrades.
### Migration path to libp2p
When the network has enough volume to justify the complexity:
1. Replace the WebSocket gossip layer with libp2p GossipSub (same topics and payload schemas, different transport)
2. Replace mDNS + relay peer list with Kademlia DHT
3. Replace circuit relay with libp2p circuit relay v2
The gossip envelope schema (`topic`, `version`, `from_peer`, `payload`) is the stable contract. As long as messages on the wire are identical, the transport layer can be swapped without touching node business logic.
## Alternatives rejected
**libp2p from the start**: `py-libp2p` is experimental and not production-ready. A Go libp2p sidecar is operationally complex. The benefits of real libp2p (mesh routing, Kademlia DHT, hole-punching) are not needed until we have hundreds of nodes.
**NATS**: Stable and fast but requires a central NATS server. Adds operational dependency and contradicts the P2P goal.
**ZeroMQ**: No NAT traversal built in. Requires manual topology management.
**No gossip (keep polling)**: Does not scale; slow to react to node churn; misses the relay/NAT requirement.

302
packages/node/meshnet_node/cli.py
View File

@@ -1,83 +1,257 @@
"""meshnet-node CLI entry point.""" """meshnet-node CLI entry point — mining-style UX."""
from __future__ import annotations
import argparse import argparse
import sys import sys
import time import time
from pathlib import Path
def _run_node(cfg: dict) -> None:
"""Start the node and hand off to the live dashboard. Blocks until Ctrl-C."""
from .startup import run_startup
from .dashboard import run_dashboard
start_time = time.monotonic()
try:
node = run_startup(
tracker_url=cfg["tracker_url"],
port=cfg.get("port", 7000),
model=cfg.get("model_name") or "stub-model",
model_id=cfg.get("model_hf_repo") or None,
shard_start=cfg.get("shard_start"),
shard_end=cfg.get("shard_end"),
quantization=cfg.get("quantization", "bfloat16").replace("bf16", "bfloat16"),
wallet_path=Path(cfg["wallet_path"]) if cfg.get("wallet_path") else None,
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg.get("host", "0.0.0.0"),
)
except Exception as exc:
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
sys.exit(1)
try:
run_dashboard(node, cfg, start_time)
except KeyboardInterrupt:
pass
finally:
node.stop()
req = getattr(node, "chat_completion_count", 0)
elapsed = time.monotonic() - start_time
h, rem = divmod(int(elapsed), 3600)
m, s = divmod(rem, 60)
print(
f"\nmeshnet-node stopped. "
f"Served {req} requests in {h:02d}:{m:02d}:{s:02d}.",
flush=True,
)
def _cmd_default(args) -> int:
"""No subcommand: wizard if no config, else start with saved config."""
from .config import load_config, save_config, merge_cli_overrides
from .wizard import run_wizard
cfg = load_config()
if cfg is None or args.reset_config:
if args.reset_config and cfg is not None:
print("Resetting config — re-running setup wizard.\n")
try:
cfg = run_wizard()
except KeyboardInterrupt:
print("\nSetup cancelled.")
return 1
save_config(cfg)
print(f"\nConfig saved to ~/.config/meshnet/config.json\n")
# Apply CLI overrides on top of saved config
overrides: dict = {}
if args.model:
overrides["model_hf_repo"] = args.model
overrides["model_name"] = args.model.split("/")[-1]
if args.quantization:
overrides["quantization"] = args.quantization
if args.download_dir:
overrides["download_dir"] = args.download_dir
if args.tracker:
overrides["tracker_url"] = args.tracker
if args.wallet:
overrides["wallet_path"] = args.wallet
if args.shard_start is not None:
overrides["shard_start"] = args.shard_start
if args.shard_end is not None:
overrides["shard_end"] = args.shard_end
if args.port is not None:
overrides["port"] = args.port
if args.host:
overrides["host"] = args.host
if overrides:
cfg = merge_cli_overrides(cfg, **overrides)
_run_node(cfg)
return 0
def _cmd_models(args) -> int:
"""List curated models (with optional HF Hub browse)."""
from .wizard import print_models_table, _browse_hf_interactive
if args.browse:
from .model_catalog import browse_hf_hub
print("Fetching HuggingFace Hub top models...\n")
try:
models = browse_hf_hub(top_n=20)
print(f"{'#':<4} {'Repo':<60} {'Downloads':>12}")
print(f"{'─'*4} {'─'*60} {'─'*12}")
for i, m in enumerate(models, 1):
dl = m["downloads"]
dl_str = (
f"{dl/1e6:.1f}M" if dl >= 1_000_000
else f"{dl/1e3:.0f}k" if dl >= 1000
else str(dl)
)
print(f"{i:<4} {m['repo']:<60} {dl_str:>12}")
except RuntimeError as exc:
print(f"Error: {exc}", file=sys.stderr)
return 1
else:
print_models_table()
return 0
def _cmd_config(args) -> int:
"""Print current config."""
import json
from .config import load_config, config_path
cfg = load_config()
if cfg is None:
print("No config file found. Run `meshnet-node` to start setup.")
return 1
print(f"Config: {config_path()}")
print(json.dumps(cfg, indent=2))
return 0
def _cmd_start(args) -> int:
"""Legacy `start` subcommand — preserves backward compatibility with existing tests."""
from .config import load_config, DEFAULTS
# Build a transient config from flags (don't write to disk)
cfg = dict(DEFAULTS)
cfg["tracker_url"] = args.tracker
cfg["port"] = args.port
cfg["model_name"] = args.model
cfg["quantization"] = args.quantization
cfg["host"] = args.host
if args.model_id:
cfg["model_hf_repo"] = args.model_id
if args.shard_start is not None:
cfg["shard_start"] = args.shard_start
if args.shard_end is not None:
cfg["shard_end"] = args.shard_end
if args.wallet:
cfg["wallet_path"] = args.wallet
if args.download_dir:
cfg["download_dir"] = args.download_dir
# Legacy start: just run without the dashboard (keep original blocking loop)
from .startup import run_startup
try:
node = run_startup(
tracker_url=cfg["tracker_url"],
port=cfg["port"],
model=cfg["model_name"],
model_id=cfg.get("model_hf_repo"),
shard_start=cfg.get("shard_start"),
shard_end=cfg.get("shard_end"),
quantization=cfg["quantization"].replace("bf16", "bfloat16"),
wallet_path=Path(cfg["wallet_path"]) if cfg.get("wallet_path") else None,
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg["host"],
advertise_host=getattr(args, "advertise_host", None),
)
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
sys.exit(1)
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
node.stop()
return 0
def main() -> None: def main() -> None:
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
prog="meshnet-node", prog="meshnet-node",
description="Distributed Inference Network node client", description="Distributed AI Inference — Node Client",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=(
"Run with no arguments to start the setup wizard.\n"
"After first setup, `meshnet-node` starts using your saved config.\n\n"
"Subcommands:\n"
" models List supported models\n"
" models --browse Browse HuggingFace Hub\n"
" config Show current config\n"
),
) )
# Flags that apply to the no-subcommand (default) path
parser.add_argument("--model", metavar="HF_REPO", help="HuggingFace repo ID to serve")
parser.add_argument("--quantization", "-q", choices=["bf16", "int8", "nf4", "bfloat16"],
help="Quantization level")
parser.add_argument("--download-dir", metavar="PATH", help="Model download directory")
parser.add_argument("--tracker", metavar="URL", help="Tracker URL")
parser.add_argument("--wallet", metavar="PATH", help="Wallet file path")
parser.add_argument("--shard-start", type=int, metavar="N", help="Pin shard start layer")
parser.add_argument("--shard-end", type=int, metavar="N", help="Pin shard end layer")
parser.add_argument("--port", type=int, metavar="N", help="Port to listen on")
parser.add_argument("--host", metavar="ADDR", help="Interface to bind (default 0.0.0.0)")
parser.add_argument("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)")
parser.add_argument("--compact", action="store_true", help="Single-line status output")
parser.add_argument("--reset-config", action="store_true", help="Re-run wizard even if config exists")
subparsers = parser.add_subparsers(dest="command") subparsers = parser.add_subparsers(dest="command")
start_cmd = subparsers.add_parser("start", help="Start the node server") # models subcommand
start_cmd.add_argument( models_cmd = subparsers.add_parser("models", help="List supported models")
"--tracker", default="http://localhost:8080", help="Tracker URL" models_cmd.add_argument("--browse", action="store_true", help="Browse HuggingFace Hub top-20")
)
start_cmd.add_argument("--port", type=int, default=7000, help="Port to listen on") # config subcommand
start_cmd.add_argument( subparsers.add_parser("config", help="Show current saved config")
"--model", default="stub-model", help="Model preset to request from tracker"
) # start subcommand (legacy / backward-compat)
start_cmd.add_argument( start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
"--model-id", start_cmd.add_argument("--tracker", default="http://localhost:8080")
help="HuggingFace model id for the real PyTorch backend", start_cmd.add_argument("--port", type=int, default=7000)
) start_cmd.add_argument("--model", default="stub-model")
start_cmd.add_argument("--shard-start", type=int, help="First layer index for an explicit shard") start_cmd.add_argument("--model-id", help="HuggingFace repo ID")
start_cmd.add_argument("--shard-end", type=int, help="Exclusive layer end index for an explicit shard") start_cmd.add_argument("--shard-start", type=int)
start_cmd.add_argument( start_cmd.add_argument("--shard-end", type=int)
"--quantization", start_cmd.add_argument("--quantization", choices=["bfloat16", "int8", "nf4", "bf16"], default="bfloat16")
choices=["bfloat16", "int8", "nf4"], start_cmd.add_argument("--host", default="0.0.0.0")
default="bfloat16", start_cmd.add_argument("--advertise-host")
help="Weight quantization for the real PyTorch backend", start_cmd.add_argument("--tracker-mode", action="store_true")
) start_cmd.add_argument("--tracker-url", default=None)
start_cmd.add_argument( start_cmd.add_argument("--wallet")
"--host", default="0.0.0.0", help="Interface to bind to" start_cmd.add_argument("--download-dir")
)
start_cmd.add_argument(
"--advertise-host",
help="Reachable host/IP to advertise to the tracker (defaults to FQDN when binding 0.0.0.0)",
)
start_cmd.add_argument(
"--tracker-mode",
action="store_true",
help="Enable client-facing /v1/chat/completions (auto-enabled when shard-start=0)",
)
start_cmd.add_argument(
"--tracker-url",
default=None,
help="Tracker URL for route selection (used in tracker mode)",
)
args = parser.parse_args() args = parser.parse_args()
if args.command == "start": if args.command == "models":
from meshnet_node.startup import run_startup sys.exit(_cmd_models(args))
elif args.command == "config":
try: sys.exit(_cmd_config(args))
node = run_startup( elif args.command == "start":
tracker_url=args.tracker, sys.exit(_cmd_start(args))
port=args.port,
model=args.model,
model_id=args.model_id,
shard_start=args.shard_start,
shard_end=args.shard_end,
quantization=args.quantization,
host=args.host,
advertise_host=args.advertise_host,
)
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
sys.exit(1)
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
node.stop()
sys.exit(0)
else: else:
parser.print_help() # Default: wizard or start with saved config
sys.exit(_cmd_default(args))
if __name__ == "__main__": if __name__ == "__main__":

72
packages/node/meshnet_node/config.py Normal file
View File

@@ -0,0 +1,72 @@
"""Persistent node configuration — stored in ~/.config/meshnet/config.json."""
from __future__ import annotations
import json
import os
import stat
from pathlib import Path
_DEFAULT_CONFIG_DIR = Path.home() / ".config" / "meshnet"
_DEFAULT_CONFIG_FILE = _DEFAULT_CONFIG_DIR / "config.json"
_DEFAULT_DOWNLOAD_DIR = Path.home() / ".meshnet" / "models"
_DEFAULT_TRACKER_URL = "http://localhost:8080"
_DEFAULT_WALLET_PATH = str(Path.home() / ".config" / "meshnet" / "wallet.json")
_DEFAULT_QUANTIZATION = "nf4"
DEFAULTS = {
"model_hf_repo": "",
"model_name": "",
"quantization": _DEFAULT_QUANTIZATION,
"download_dir": str(_DEFAULT_DOWNLOAD_DIR),
"tracker_url": _DEFAULT_TRACKER_URL,
"wallet_path": _DEFAULT_WALLET_PATH,
"shard_start": None,
"shard_end": None,
"port": 7000,
"host": "0.0.0.0",
}
def config_path(override: Path | None = None) -> Path:
return override or _DEFAULT_CONFIG_FILE
def load_config(path: Path | None = None) -> dict | None:
"""Return parsed config dict, or None if no config file exists."""
p = config_path(path)
if not p.exists():
return None
try:
cfg = json.loads(p.read_text())
if not isinstance(cfg, dict):
return None
return cfg
except (json.JSONDecodeError, OSError):
return None
def save_config(cfg: dict, path: Path | None = None) -> None:
"""Write config to disk with restricted permissions (0o600)."""
p = config_path(path)
p.parent.mkdir(parents=True, exist_ok=True)
p.write_text(json.dumps(cfg, indent=2))
try:
os.chmod(p, stat.S_IRUSR | stat.S_IWUSR)
except OSError:
pass # Windows / some filesystems don't support chmod
def delete_config(path: Path | None = None) -> None:
p = config_path(path)
if p.exists():
p.unlink()
def merge_cli_overrides(cfg: dict, **cli_kwargs) -> dict:
"""Return a copy of cfg with any non-None CLI values applied on top."""
result = dict(cfg)
for key, val in cli_kwargs.items():
if val is not None:
result[key] = val
return result

220
packages/node/meshnet_node/dashboard.py Normal file
View File

@@ -0,0 +1,220 @@
"""Live node status dashboard — rich TUI with plain-text fallback."""
from __future__ import annotations
import os
import sys
import time
from collections import deque
from typing import TYPE_CHECKING
if TYPE_CHECKING:
pass
def is_interactive_tty() -> bool:
"""Return True when stdout is a real terminal (not CI / redirected / WSL2 dumb)."""
if not sys.stdout.isatty():
return False
term = os.environ.get("TERM", "")
if term in ("dumb", ""):
return False
return True
def _format_uptime(seconds: float) -> str:
s = int(seconds)
h, rem = divmod(s, 3600)
m, sec = divmod(rem, 60)
return f"{h:02d}:{m:02d}:{sec:02d}"
def _gpu_stats() -> list[dict]:
"""Return per-GPU utilization and VRAM stats, or empty list on CPU."""
try:
import torch # type: ignore[import]
if not torch.cuda.is_available():
return []
stats = []
for i in range(torch.cuda.device_count()):
props = torch.cuda.get_device_properties(i)
used = torch.cuda.memory_allocated(i)
total = props.total_memory
# Utilization requires pynvml; skip gracefully if not available
util = _nvml_gpu_util(i)
stats.append(
{
"index": i,
"name": props.name,
"used_gb": used / 1e9,
"total_gb": total / 1e9,
"util_pct": util,
}
)
return stats
except ImportError:
return []
def _nvml_gpu_util(index: int) -> int | None:
"""Return GPU utilization % via pynvml, or None if unavailable."""
try:
import pynvml # type: ignore[import]
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
rates = pynvml.nvmlDeviceGetUtilizationRates(handle)
return rates.gpu
except Exception:
return None
class _EMA:
"""Exponential moving average for tokens/sec."""
def __init__(self, alpha: float = 0.1):
self._alpha = alpha
self._value: float | None = None
def update(self, sample: float) -> float:
if self._value is None:
self._value = sample
else:
self._value = self._alpha * sample + (1 - self._alpha) * self._value
return self._value
@property
def value(self) -> float:
return self._value or 0.0
def _make_bar(pct: float, width: int = 10) -> str:
filled = round(pct / 100 * width)
return "█" * filled + "░" * (width - filled)
def run_dashboard(node, config: dict, start_time: float) -> None:
"""Start the live dashboard. Blocks until Ctrl-C. Returns cleanly."""
if not is_interactive_tty():
_run_plain_loop(node, config, start_time)
return
try:
from rich.live import Live # type: ignore[import]
_run_rich_dashboard(node, config, start_time)
except ImportError:
_run_plain_loop(node, config, start_time)
def _build_rich_renderable(
node, config: dict, start_time: float, tps_ema: _EMA, prev_req: list[int]
):
from rich.table import Table # type: ignore[import]
from rich.panel import Panel # type: ignore[import]
from rich.columns import Columns # type: ignore[import]
from rich.text import Text # type: ignore[import]
uptime = time.monotonic() - start_time
req_count = getattr(node, "chat_completion_count", 0)
# Tokens/sec EMA (approximate: 20 tokens per request heuristic when no real counter)
delta_req = req_count - prev_req[0]
prev_req[0] = req_count
if delta_req > 0:
approx_tokens = delta_req * 20
tps_ema.update(approx_tokens / 2.0) # 2s interval
gpu_stats = _gpu_stats()
model_name = config.get("model_name") or config.get("model_hf_repo", "unknown").split("/")[-1]
shard = ""
if config.get("shard_start") is not None:
shard = f" shard {config['shard_start']}–{config['shard_end']}"
# Header line
header = Text(
f"meshnet-node {model_name} [{config.get('quantization', 'bf16')}]{shard}"
f" up {_format_uptime(uptime)}",
style="bold white",
)
# GPU table
gpu_table = Table(show_header=False, box=None, padding=(0, 1))
gpu_table.add_column("label", style="dim", no_wrap=True)
gpu_table.add_column("bar", no_wrap=True)
gpu_table.add_column("vram", no_wrap=True, style="cyan")
if gpu_stats:
for g in gpu_stats:
util = g["util_pct"]
util_str = f"{_make_bar(util)} {util:3d}%" if util is not None else " n/a"
vram_str = f"VRAM {g['used_gb']:.1f}/{g['total_gb']:.1f} GB"
gpu_table.add_row(f"GPU {g['index']} {g['name'][:20]}", util_str, vram_str)
else:
gpu_table.add_row("CPU mode", "", "no GPU detected")
# Stats panel
tps = tps_ema.value
bar_len = min(8, max(0, int(tps / 10)))
tps_bar = "▁▂▃▄▅▆▇█"[:bar_len].ljust(8)
stats_lines = [
f"Tokens/sec {tps_bar} {tps:.1f} t/s (EMA)",
f"Requests {req_count:,} served",
f"Peers 0 connected (gossip: US-017)",
f"TAI earned 0.00 TAI (payments: US-006)",
f"Uptime {_format_uptime(uptime)}",
"",
"[q] quit [c] compact view",
]
from rich.console import Group # type: ignore[import]
return Panel(
Group(header, gpu_table, Text("\n".join(stats_lines))),
title="[bold green]meshnet-node[/bold green]",
border_style="green",
)
def _run_rich_dashboard(node, config: dict, start_time: float) -> None:
from rich.live import Live # type: ignore[import]
tps_ema = _EMA()
prev_req = [0]
try:
with Live(
_build_rich_renderable(node, config, start_time, tps_ema, prev_req),
refresh_per_second=0.5,
screen=False,
) as live:
while True:
time.sleep(2)
live.update(
_build_rich_renderable(node, config, start_time, tps_ema, prev_req)
)
except KeyboardInterrupt:
pass
def _run_plain_loop(node, config: dict, start_time: float) -> None:
model_name = config.get("model_name") or config.get("model_hf_repo", "unknown").split("/")[-1]
try:
while True:
uptime = time.monotonic() - start_time
req = getattr(node, "chat_completion_count", 0)
gpu_stats = _gpu_stats()
vram_str = ""
if gpu_stats:
g = gpu_stats[0]
vram_str = f" VRAM{g['used_gb']:.1f}GB"
print(
f"[{model_name} req{req}{vram_str} up{_format_uptime(uptime)}]",
flush=True,
)
time.sleep(2)
except KeyboardInterrupt:
pass

18
packages/node/meshnet_node/model_backend.py
View File

@@ -241,6 +241,24 @@ class TorchModelShard:
yield token_text yield token_text
t.join() t.join()
def count_prompt_tokens(self, messages: list[dict]) -> int:
"""Return tokenizer-backed prompt token count for OpenAI usage metadata."""
encoded = self._encode_messages(messages)
input_ids = encoded["input_ids"]
return int(input_ids.shape[-1])
def count_text_tokens(self, text: str) -> int:
"""Return tokenizer-backed completion token count for OpenAI usage metadata."""
try:
encoded = self.tokenizer(
text,
return_tensors="pt",
add_special_tokens=False,
)
except TypeError:
encoded = self.tokenizer(text, return_tensors="pt")
return int(encoded["input_ids"].shape[-1])
def _encode_messages(self, messages: list[dict]) -> dict: def _encode_messages(self, messages: list[dict]) -> dict:
"""Format messages with chat template (if available) and tokenize.""" """Format messages with chat template (if available) and tokenize."""
if hasattr(self.tokenizer, "apply_chat_template"): if hasattr(self.tokenizer, "apply_chat_template"):

165
packages/node/meshnet_node/model_catalog.py Normal file
View File

@@ -0,0 +1,165 @@
"""Curated list of models supported by the network with VRAM requirements."""
from __future__ import annotations
from dataclasses import dataclass
@dataclass
class ModelPreset:
name: str
hf_repo: str
num_layers: int
# VRAM in GB at each quantization level (None = too large to quantize this way)
vram_nf4: float
vram_int8: float
vram_bf16: float
description: str
def vram_for_quant(self, quant: str) -> float:
"""Return VRAM requirement in GB for the given quantization."""
q = quant.lower().replace("bfloat16", "bf16")
if q == "nf4":
return self.vram_nf4
if q in ("int8", "int8"):
return self.vram_int8
if q in ("bf16", "bfloat16"):
return self.vram_bf16
raise ValueError(f"unknown quantization: {quant!r}")
def fits_vram(self, available_gb: float, quant: str) -> bool:
return self.vram_for_quant(quant) <= available_gb
def recommended_quant(self, available_gb: float) -> str | None:
"""Return the highest-quality quantization that fits available VRAM, or None."""
if self.vram_bf16 <= available_gb:
return "bf16"
if self.vram_int8 <= available_gb:
return "int8"
if self.vram_nf4 <= available_gb:
return "nf4"
return None
CURATED_MODELS: list[ModelPreset] = [
ModelPreset(
name="Qwen2.5-0.5B-Instruct",
hf_repo="Qwen/Qwen2.5-0.5B-Instruct",
num_layers=24,
vram_nf4=0.4,
vram_int8=0.6,
vram_bf16=1.0,
description="Smallest no-gating model — great for testing, ~1 GB",
),
ModelPreset(
name="Qwen2.5-1.5B-Instruct",
hf_repo="Qwen/Qwen2.5-1.5B-Instruct",
num_layers=28,
vram_nf4=1.0,
vram_int8=1.8,
vram_bf16=3.2,
description="Fast no-gating model — good quality, ~3 GB",
),
ModelPreset(
name="Llama-3-70B-Instruct",
hf_repo="meta-llama/Meta-Llama-3-70B-Instruct",
num_layers=80,
vram_nf4=18.0,
vram_int8=40.0,
vram_bf16=140.0,
description="Meta's flagship 70B instruction model",
),
ModelPreset(
name="Qwen2.5-72B-Instruct",
hf_repo="Qwen/Qwen2.5-72B-Instruct",
num_layers=80,
vram_nf4=19.0,
vram_int8=41.0,
vram_bf16=145.0,
description="Alibaba's 72B multilingual instruction model",
),
ModelPreset(
name="Mixtral-8x7B-Instruct",
hf_repo="mistralai/Mixtral-8x7B-Instruct-v0.1",
num_layers=32,
vram_nf4=7.0,
vram_int8=14.0,
vram_bf16=27.0,
description="Mistral's sparse MoE — fast and efficient",
),
ModelPreset(
name="Llama-3-8B-Instruct",
hf_repo="meta-llama/Meta-Llama-3-8B-Instruct",
num_layers=32, # gated repo — requires HF login
vram_nf4=4.5,
vram_int8=8.5,
vram_bf16=16.0,
description="Meta's compact 8B model — good for low-VRAM nodes",
),
ModelPreset(
name="Phi-3-medium-128k",
hf_repo="microsoft/Phi-3-medium-128k-instruct",
num_layers=40,
vram_nf4=4.0,
vram_int8=8.0,
vram_bf16=15.0,
description="Microsoft's efficient 14B model with 128k context",
),
ModelPreset(
name="Gemma-2-27B-IT",
hf_repo="google/gemma-2-27b-it",
num_layers=46,
vram_nf4=10.0,
vram_int8=20.0,
vram_bf16=54.0,
description="Google's 27B instruction-tuned model",
),
ModelPreset(
name="DeepSeek-V2-Lite-Chat",
hf_repo="deepseek-ai/DeepSeek-V2-Lite-Chat",
num_layers=27,
vram_nf4=5.0,
vram_int8=9.0,
vram_bf16=16.0,
description="DeepSeek's efficient MoE — strong coding + reasoning",
),
]
def detect_num_layers(hf_repo: str) -> int | None:
"""Return num_hidden_layers from HuggingFace config.json (downloads ~1 KB only)."""
# Check curated list first (no network call)
for m in CURATED_MODELS:
if m.hf_repo == hf_repo:
return m.num_layers
try:
from transformers import AutoConfig # type: ignore[import]
cfg = AutoConfig.from_pretrained(hf_repo)
return int(cfg.num_hidden_layers)
except Exception:
return None
def browse_hf_hub(top_n: int = 20) -> list[dict]:
"""Fetch top downloaded text-generation models from HuggingFace Hub."""
try:
from huggingface_hub import list_models # type: ignore[import]
models = list(
list_models(
pipeline_tag="text-generation",
library="transformers",
sort="downloads",
direction=-1,
limit=top_n,
)
)
return [
{
"repo": m.id,
"downloads": getattr(m, "downloads", 0) or 0,
}
for m in models
]
except Exception as exc:
raise RuntimeError(f"HuggingFace Hub lookup failed: {exc}") from exc

38
packages/node/meshnet_node/startup.py
View File

@@ -84,7 +84,19 @@ def run_startup(
if probationary_line is not None: if probationary_line is not None:
print(f" {probationary_line}", flush=True) print(f" {probationary_line}", flush=True)
if model_id is not None and shard_start is not None and shard_end is not None: if model_id is not None:
# Auto-detect shard range from model config if not explicitly provided
if shard_start is None or shard_end is None:
detected = _detect_num_layers(model_id)
if detected is None:
raise ValueError(
f"Could not read num_hidden_layers from {model_id} config. "
"Pass --shard-start and --shard-end explicitly."
)
shard_start = shard_start if shard_start is not None else 0
shard_end = shard_end if shard_end is not None else detected - 1
print(f" Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True)
print("Loading real PyTorch model shard...", flush=True) print("Loading real PyTorch model shard...", flush=True)
node = TorchNodeServer( node = TorchNodeServer(
host=host, host=host,
@@ -93,8 +105,15 @@ def run_startup(
shard_start=shard_start, shard_start=shard_start,
shard_end=shard_end, shard_end=shard_end,
quantization=quantization, quantization=quantization,
tracker_url=tracker_url,
) )
actual_port = node.start() actual_port = node.start()
total_layers = getattr(node.backend, "total_layers", None)
if isinstance(total_layers, int) and total_layers > 0:
layer_count = shard_end - shard_start + 1
shard_label = f"layers {shard_start}–{shard_end}; {layer_count} of {total_layers}"
else:
shard_label = f"layers {shard_start}–{shard_end}"
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}" endpoint = f"http://{public_host}:{actual_port}"
print( print(
@@ -102,7 +121,7 @@ def run_startup(
f"meshnet-node ready\n" f"meshnet-node ready\n"
f" Wallet: {address}\n" f" Wallet: {address}\n"
f" Model ID: {model_id}\n" f" Model ID: {model_id}\n"
f" Shard: layers {shard_start}-{shard_end}\n" f" Shard: {shard_label}\n"
f" Quantization: {quantization}\n" f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n" f" Endpoint: {endpoint}\n"
f" Hardware: {device.upper()}\n" f" Hardware: {device.upper()}\n"
@@ -110,8 +129,8 @@ def run_startup(
flush=True, flush=True,
) )
return node return node
if model_id is not None or shard_start is not None or shard_end is not None: if shard_start is not None or shard_end is not None:
raise ValueError("--model-id, --shard-start, and --shard-end must be provided together") raise ValueError("--shard-start / --shard-end require --model-id")
# 3. Shard assignment from tracker # 3. Shard assignment from tracker
print("Querying tracker for shard assignment...", flush=True) print("Querying tracker for shard assignment...", flush=True)
@@ -201,6 +220,17 @@ def run_startup(
return node return node
def _detect_num_layers(model_id: str) -> int | None:
"""Fetch num_hidden_layers from HuggingFace model config (downloads ~1 KB config.json only)."""
try:
from transformers import AutoConfig # type: ignore[import]
cfg = AutoConfig.from_pretrained(model_id)
return int(cfg.num_hidden_layers)
except Exception as exc:
print(f" Warning: could not read model config from HF: {exc}", flush=True)
return None
def _probationary_status_line(contracts: Any | None, wallet_address: str) -> str | None: def _probationary_status_line(contracts: Any | None, wallet_address: str) -> str | None:
if contracts is None: if contracts is None:
return None return None

62
packages/node/meshnet_node/torch_server.py
View File

@@ -11,6 +11,7 @@ import urllib.error
import urllib.parse import urllib.parse
import urllib.request import urllib.request
import uuid import uuid
from typing import Any
from .model_backend import ( from .model_backend import (
InsufficientVRAMError, InsufficientVRAMError,
@@ -232,7 +233,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
) )
else: else:
text = server.backend.generate_text(messages, max_tokens, temperature, top_p) text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
self._send_openai_response(text, model_name, False) self._send_openai_response(text, model_name, False, messages)
except Exception as exc: except Exception as exc:
self._send_json(500, {"error": f"generation failed: {exc}"}) self._send_json(500, {"error": f"generation failed: {exc}"})
return return
@@ -250,7 +251,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
return return
remaining_route = self._get_remaining_route(model_name) remaining_route = self._get_remaining_route(model_name)
result_text = self._run_downstream_pipeline(payload, remaining_route) result_text = self._run_downstream_pipeline(payload, remaining_route)
self._send_openai_response(result_text, model_name, stream) self._send_openai_response(result_text, model_name, stream, messages)
def _get_remaining_route(self, model: str) -> list[str]: def _get_remaining_route(self, model: str) -> list[str]:
server: _TorchHTTPServer = self.server # type: ignore[assignment] server: _TorchHTTPServer = self.server # type: ignore[assignment]
@@ -367,10 +368,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.wfile.write(b"data: [DONE]\n\n") self.wfile.write(b"data: [DONE]\n\n")
self.wfile.flush() self.wfile.flush()
def _send_openai_response(self, text: str, model: str, stream: bool) -> None: def _send_openai_response(
self,
text: str,
model: str,
stream: bool,
messages: list[dict] | None = None,
) -> None:
chunk_id = "chatcmpl-node" chunk_id = "chatcmpl-node"
created = int(time.time()) created = int(time.time())
if not stream: if not stream:
usage = _usage_for_response(self.server.backend, messages or [], text) # type: ignore[attr-defined]
self._send_json(200, { self._send_json(200, {
"id": chunk_id, "id": chunk_id,
"object": "chat.completion", "object": "chat.completion",
@@ -381,7 +389,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"message": {"role": "assistant", "content": text}, "message": {"role": "assistant", "content": text},
"finish_reason": "stop", "finish_reason": "stop",
}], }],
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, "usage": usage,
}) })
return return
self.send_response(200) self.send_response(200)
@@ -412,6 +420,52 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.wfile.flush() self.wfile.flush()
def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]:
prompt_tokens = _backend_token_count(
backend,
"count_prompt_tokens",
messages,
fallback=_fallback_message_token_count(messages),
)
completion_tokens = _backend_token_count(
backend,
"count_text_tokens",
completion_text,
fallback=_fallback_text_token_count(completion_text),
)
return {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
}
def _backend_token_count(backend: object, method_name: str, value: object, fallback: int) -> int:
method: Any = getattr(backend, method_name, None)
if callable(method):
try:
return max(0, int(method(value)))
except Exception:
pass
return max(0, int(fallback))
def _fallback_message_token_count(messages: list[dict]) -> int:
text = " ".join(
str(message.get("content", ""))
for message in messages
if isinstance(message, dict)
)
return _fallback_text_token_count(text)
def _fallback_text_token_count(text: str) -> int:
parts = text.split()
if parts:
return len(parts)
return 1 if text else 0
class TorchNodeServer: class TorchNodeServer:
"""HTTP server backed by a HuggingFace causal language model shard.""" """HTTP server backed by a HuggingFace causal language model shard."""

332
packages/node/meshnet_node/wizard.py Normal file
View File

@@ -0,0 +1,332 @@
"""Interactive first-run setup wizard — mining-client style."""
from __future__ import annotations
import sys
import urllib.error
import urllib.request
from pathlib import Path
from typing import TYPE_CHECKING
from .config import DEFAULTS, _DEFAULT_DOWNLOAD_DIR, _DEFAULT_TRACKER_URL, _DEFAULT_WALLET_PATH
from .model_catalog import CURATED_MODELS, ModelPreset, browse_hf_hub, detect_num_layers
if TYPE_CHECKING:
pass
_HEADER = """\
╔══════════════════════════════════════════════════════════════════╗
║ meshnet-node v0.1.0 ║
║ Distributed AI Inference — Node Setup ║
╚══════════════════════════════════════════════════════════════════╝
"""
_QUANT_LABELS = {"nf4": "NF4 (4-bit)", "int8": "INT8 (8-bit)", "bf16": "BF16 (full)"}
def _ask(prompt: str, default: str = "", validator=None) -> str:
"""Prompt user and return answer. Returns default on empty input or EOF."""
display = f"{prompt} [{default}]: " if default else f"{prompt}: "
while True:
try:
raw = input(display).strip()
except (EOFError, KeyboardInterrupt):
print()
raise KeyboardInterrupt
value = raw or default
if validator is None or validator(value):
return value
# validator returned error string
print(f" ✗ {validator(value)}")
def _ask_int(prompt: str, default: int, lo: int, hi: int) -> int:
def validate(s: str) -> bool | str:
try:
v = int(s)
except ValueError:
return "Please enter a number."
if not (lo <= v <= hi):
return f"Please enter a number between {lo} and {hi}."
return True
while True:
raw = _ask(prompt, str(default))
try:
v = int(raw)
if lo <= v <= hi:
return v
except ValueError:
pass
print(f" ✗ Enter a number between {lo} and {hi}.")
def _ask_yn(prompt: str, default: bool = True) -> bool:
hint = "Y/n" if default else "y/N"
raw = _ask(f"{prompt} [{hint}]").lower()
if not raw:
return default
return raw.startswith("y")
def _detect_gpus() -> list[dict]:
"""Return list of detected GPU dicts with name and vram_gb."""
gpus: list[dict] = []
try:
import torch # type: ignore[import]
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
props = torch.cuda.get_device_properties(i)
gpus.append(
{
"index": i,
"name": props.name,
"vram_gb": props.total_memory / 1e9,
"backend": "cuda",
}
)
except ImportError:
pass
return gpus
def _total_vram_gb(gpus: list[dict]) -> float:
return sum(g["vram_gb"] for g in gpus)
def _print_gpus(gpus: list[dict]) -> None:
if not gpus:
print(" ⚠ No CUDA GPU detected — running in CPU mode")
print(" CPU inference is very slow. Consider a machine with an NVIDIA GPU.")
return
for g in gpus:
vram = g["vram_gb"]
print(f" GPU {g['index']}: {g['name']} {vram:.0f} GB VRAM ✓")
def _print_model_table(gpus: list[dict], quant: str = "nf4") -> None:
available_gb = _total_vram_gb(gpus)
print()
print(f" # {'Model':<30} {'Layers':>6} {'NF4':>6} {'INT8':>6} {'BF16':>6}")
print(f" {'─'*4} {'─'*30} {'─'*6} {'─'*6} {'─'*6} {'─'*6}")
for i, m in enumerate(CURATED_MODELS, 1):
fits_nf4 = "✓" if m.vram_nf4 <= available_gb else "✗"
fits_int8 = "✓" if m.vram_int8 <= available_gb else "✗"
fits_bf16 = "✓" if m.vram_bf16 <= available_gb else "✗"
nf4_str = f"{fits_nf4}{m.vram_nf4:.0f}GB"
int8_str = f"{fits_int8}{m.vram_int8:.0f}GB"
bf16_str = f"{fits_bf16}{m.vram_bf16:.0f}GB"
print(f" {i:<3} {m.name:<30} {m.num_layers:>6} {nf4_str:>6} {int8_str:>6} {bf16_str:>6}")
print(f" {m.description}")
idx = len(CURATED_MODELS) + 1
print(f" {idx:<3} {'[Browse HuggingFace Hub...]':<30}")
print()
def _browse_hf_interactive() -> str | None:
"""Show HF Hub top-20 and let user enter a repo ID. Returns repo ID or None to go back."""
print("\nFetching top models from HuggingFace Hub...")
try:
models = browse_hf_hub(top_n=20)
except RuntimeError as exc:
print(f" ✗ {exc}")
return None
print(f"\n {'#':<4} {'HuggingFace Repo':<50} Downloads")
print(f" {'─'*4} {'─'*50} {'─'*10}")
for i, m in enumerate(models, 1):
dl = m["downloads"]
dl_str = f"{dl/1e6:.1f}M" if dl >= 1_000_000 else f"{dl/1e3:.0f}k" if dl >= 1000 else str(dl)
print(f" {i:<4} {m['repo']:<50} {dl_str}")
print()
raw = _ask(
"Enter a number to select, or paste any HuggingFace repo ID (or press Enter to go back)",
default="",
)
if not raw:
return None
try:
idx = int(raw) - 1
if 0 <= idx < len(models):
return models[idx]["repo"]
except ValueError:
pass
# Treat raw input as a repo ID
if "/" in raw:
return raw
print(" ✗ Invalid input — please enter a number or a full repo ID like 'org/model-name'")
return None
def _ask_quant(gpus: list[dict], model: ModelPreset | None) -> str:
available_gb = _total_vram_gb(gpus)
print("\nQuantization level:")
options: list[tuple[str, str]] = []
for quant, label in [("nf4", "NF4 4-bit"), ("int8", "INT8 8-bit"), ("bf16", "BF16 full precision")]:
if model is not None:
vram = model.vram_for_quant(quant)
fits = "✓" if vram <= available_gb else "✗ insufficient VRAM"
suffix = f" ({vram:.0f} GB needed — {fits})"
else:
suffix = ""
options.append((quant, f"{label}{suffix}"))
for i, (_, label) in enumerate(options, 1):
print(f" {i}) {label}")
# Recommend the best fitting quant
if model is not None:
rec = model.recommended_quant(available_gb)
rec_idx = next((i for i, (q, _) in enumerate(options, 1) if q == rec), 1) if rec else 1
default_idx = rec_idx
print(f" (Recommended: {rec.upper() if rec else 'NF4'} for your GPU)")
else:
default_idx = 1
choice = _ask_int("Enter number", default_idx, 1, 3)
return options[choice - 1][0]
def _validate_dir(path_str: str) -> bool | str:
p = Path(path_str).expanduser()
try:
p.mkdir(parents=True, exist_ok=True)
return True
except OSError as exc:
return f"Cannot create directory: {exc}"
def _validate_tracker(url: str) -> bool | str:
if not url.startswith(("http://", "https://")):
return "URL must start with http:// or https://"
return True
def _ping_tracker(url: str) -> bool:
"""Return True if tracker responds to /health."""
try:
with urllib.request.urlopen(f"{url.rstrip('/')}/health", timeout=3):
return True
except Exception:
return False
def run_wizard(config_path_override=None) -> dict:
"""Run the interactive setup wizard and return a config dict.
Raises KeyboardInterrupt if user presses Ctrl-C.
"""
print(_HEADER)
# Step 1: GPU detection
print("Detecting hardware...")
gpus = _detect_gpus()
_print_gpus(gpus)
available_gb = _total_vram_gb(gpus)
if available_gb == 0:
available_gb = 9999 # CPU — don't filter models by VRAM
# Step 2 & 3: Model selection
print("\nSelect a model to serve:\n")
selected_repo: str | None = None
selected_preset: ModelPreset | None = None
while selected_repo is None:
_print_model_table(gpus)
lo, hi = 1, len(CURATED_MODELS) + 1
choice = _ask_int("Enter number", 1, lo, hi)
if choice == len(CURATED_MODELS) + 1:
repo = _browse_hf_interactive()
if repo:
# Look up layer count for custom repo
print(f" Checking {repo} config...", end=" ", flush=True)
layers = detect_num_layers(repo)
if layers:
print(f"{layers} layers")
else:
print("(layer count unknown — will detect on start)")
selected_repo = repo
selected_preset = None
else:
selected_preset = CURATED_MODELS[choice - 1]
selected_repo = selected_preset.hf_repo
if selected_preset.recommended_quant(available_gb) is None:
print(
f"\n ⚠ Warning: {selected_preset.name} requires at least "
f"{selected_preset.vram_nf4:.0f} GB VRAM at NF4 — even the smallest "
f"quantization may be too large for your GPU."
)
if not _ask_yn("Continue anyway?", default=False):
selected_repo = None
selected_preset = None
num_layers = (selected_preset.num_layers if selected_preset
else detect_num_layers(selected_repo or ""))
layers_str = f" {num_layers} layers" if num_layers else ""
print(f"\n ✓ Selected: {selected_repo}{layers_str}")
# Step 3b: Quantization
quant = _ask_quant(gpus, selected_preset)
print(f" ✓ Quantization: {quant.upper()}")
# Step 4: Download directory
print()
dl_dir = _ask(
"Download directory",
default=str(_DEFAULT_DOWNLOAD_DIR),
validator=lambda v: _validate_dir(v) if v else "Directory is required.",
)
print(f" ✓ Download dir: {dl_dir}")
# Step 5: Tracker URL
print()
tracker_url = _DEFAULT_TRACKER_URL
raw_tracker = _ask("Tracker URL", default=_DEFAULT_TRACKER_URL, validator=_validate_tracker)
tracker_url = raw_tracker
if _ping_tracker(tracker_url):
print(f" ✓ Tracker reachable: {tracker_url}")
else:
print(f" ⚠ Tracker not reachable at {tracker_url} (will retry on start)")
# Step 6: Wallet path
print()
wallet_path = _ask("Wallet path", default=_DEFAULT_WALLET_PATH)
print(f" ✓ Wallet: {wallet_path}")
cfg = {
"model_hf_repo": selected_repo,
"model_name": selected_preset.name if selected_preset else selected_repo.split("/")[-1],
"quantization": quant,
"download_dir": dl_dir,
"tracker_url": tracker_url,
"wallet_path": wallet_path,
"shard_start": None,
"shard_end": None,
"port": DEFAULTS["port"],
"host": DEFAULTS["host"],
}
return cfg
def print_models_table(available_gb: float | None = None) -> None:
"""Print curated model table for `meshnet-node models`."""
gpus: list[dict] = []
if available_gb is None:
gpus = _detect_gpus()
available_gb = _total_vram_gb(gpus) or 9999
else:
gpus = [{"index": 0, "name": "GPU", "vram_gb": available_gb, "backend": "cuda"}]
print(f"\n{'#':<4} {'Model':<32} {'HuggingFace Repo':<45} {'Layers':>6} {'NF4':>8} {'INT8':>8} {'BF16':>8}")
print(f"{'─'*4} {'─'*32} {'─'*45} {'─'*6} {'─'*8} {'─'*8} {'─'*8}")
for i, m in enumerate(CURATED_MODELS, 1):
def _cell(vram: float) -> str:
fits = "✓" if vram <= available_gb else "✗"
return f"{fits}{vram:.0f}GB"
print(
f"{i:<4} {m.name:<32} {m.hf_repo:<45} {m.num_layers:>6} "
f"{_cell(m.vram_nf4):>8} {_cell(m.vram_int8):>8} {_cell(m.vram_bf16):>8}"
)
print()

2
scripts/ralph_progress.py
View File

@@ -152,7 +152,7 @@ def _story_meta(
parts.append(f"worktree: {wt}") parts.append(f"worktree: {wt}")
# summary ------------------------------------------------------------- # summary -------------------------------------------------------------
notes = story.get("completionNotes", "").strip() notes = (story.get("completionNotes") or "").strip()
if not notes and wt: if not notes and wt:
notes = _story_last_commit(sid) notes = _story_last_commit(sid)
if notes: if notes:

37
tests/test_node_startup.py
View File

@@ -347,6 +347,43 @@ def test_tracker_assign_lists_peers_for_same_model_shard():
tracker.stop() tracker.stop()
def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, capsys):
"""Real-model startup summary prints the shard range plus total model layers."""
import meshnet_node.startup as startup_mod
class FakeBackend:
total_layers = 24
class FakeTorchNodeServer:
def __init__(self, **kwargs):
self.kwargs = kwargs
self.backend = FakeBackend()
self.port = None
def start(self):
self.port = 8001
return self.port
monkeypatch.setattr(
startup_mod,
"detect_hardware",
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
)
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
node = run_startup(
tracker_url="http://127.0.0.1:8080",
model_id="Qwen/Qwen2.5-0.5B-Instruct",
shard_start=0,
shard_end=23,
wallet_path=tmp_path / "wallet.json",
)
assert node.backend.total_layers == 24
output = capsys.readouterr().out
assert "Shard: layers 0–23; 24 of 24" in output
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Full startup integration test # Full startup integration test
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------

88
tests/test_real_model_backend.py
View File

@@ -12,6 +12,9 @@ import pytest
from meshnet_node.model_backend import ( from meshnet_node.model_backend import (
InsufficientVRAMError, InsufficientVRAMError,
TensorPayload, TensorPayload,
_call_layer,
_decoder_attention_mask,
_int_tensor_header,
build_quantization_config, build_quantization_config,
validate_quantization, validate_quantization,
) )
@@ -52,6 +55,32 @@ class _FakeTailBackend(_FakeBackend):
return " Paris" return " Paris"
class _FakeFullBackend(_FakeBackend):
is_head = True
is_tail = True
def generate_text(
self,
messages: list[dict],
max_new_tokens: int = 16,
temperature: float = 1.0,
top_p: float = 1.0,
) -> str:
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
assert max_new_tokens == 7
assert temperature == 1.0
assert top_p == 1.0
return "56"
def count_prompt_tokens(self, messages: list[dict]) -> int:
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
return 8
def count_text_tokens(self, text: str) -> int:
assert text == "56"
return 1
def test_quantization_flag_validation(): def test_quantization_flag_validation():
assert validate_quantization("bfloat16") == "bfloat16" assert validate_quantization("bfloat16") == "bfloat16"
assert validate_quantization("int8") == "int8" assert validate_quantization("int8") == "int8"
@@ -145,6 +174,65 @@ def test_tail_forward_returns_text_completion_from_binary_activations():
node.stop() node.stop()
def test_full_model_chat_completion_uses_generation_not_single_token_decode():
node = TorchNodeServer(backend=_FakeFullBackend())
port = node.start()
try:
payload = json.dumps({
"model": "fake-model",
"messages": [{"role": "user", "content": "What is 7 times 8?"}],
"max_tokens": 7,
}).encode()
req = urllib.request.Request(
f"http://127.0.0.1:{port}/v1/chat/completions",
data=payload,
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req, timeout=5) as resp:
body = json.loads(resp.read())
assert body["choices"][0]["message"]["content"] == "56"
assert body["usage"] == {"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9}
finally:
node.stop()
def test_int_tensor_header_serializes_torch_tensors():
torch = pytest.importorskip("torch")
header = _int_tensor_header(torch.tensor([[1, 2, 3]], dtype=torch.long))
assert header.startswith("1,3:")
def test_decoder_attention_mask_is_causal_float_mask():
torch = pytest.importorskip("torch")
hidden_states = torch.zeros((1, 3, 8), dtype=torch.bfloat16)
mask = _decoder_attention_mask(torch.ones((1, 3), dtype=torch.long), hidden_states, torch)
assert mask.shape == (1, 1, 3, 3)
assert mask.dtype == torch.bfloat16
assert mask[0, 0, 0, 1] < 0
assert mask[0, 0, 2, 0] == 0
def test_call_layer_passes_rotary_position_embeddings():
class NeedsPositionEmbeddings:
def __call__(self, hidden_states, **kwargs):
assert kwargs["position_embeddings"] == "rotary"
return hidden_states
assert _call_layer(
NeedsPositionEmbeddings(),
"hidden",
attention_mask=None,
position_ids="positions",
position_embeddings="rotary",
) == "hidden"
@pytest.mark.integration @pytest.mark.integration
def test_two_node_gpt2_completion_is_deterministic(): def test_two_node_gpt2_completion_is_deterministic():
if os.environ.get("CI"): if os.environ.get("CI"):