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

feat: add real PyTorch model backend

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-29 15:54:40 +03:00
parent c358798627
commit 2690d9b9ba
8 changed files with 877 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
.scratch/distributed-inference-network/prd.json
View File

@@ -5,8 +5,8 @@
"userStories": [
{
"id": "US-001",
"title": "01 \u2014 Monorepo scaffold + single-node smoke test",
"description": "Stand up the monorepo package layout and prove that a client HTTP request can travel end-to-end through the stack \u2014 gateway \u2192 one node serving all layers \u2192 valid response \u2014 before any real model or distributed logic is wired in. The six top-level packages should exist with minimal stubs: - `packages/node` \u2014 node client CLI (`meshnet-node`) - `packages/gateway` \u2014 HTTP gateway + route orchestration - `packages/tracker` \u2014 node registry and route selection - `packages/sdk` \u2014 `meshnet` Python SDK - `packages/contracts` \u2014 Solana smart contract wrappers - `packages/p2p` \u2014 gossip and shard swarm The smoke test uses a tiny in-process stub model (not a real LLM) so the test is fast and has no external dependencies. The node runs in the same process as the test. The gateway is a real HTTP server. The client sends a real `POST /v1/chat/completions` request and receives a valid OpenAI-format response.",
"title": "01 — Monorepo scaffold + single-node smoke test",
"description": "Stand up the monorepo package layout and prove that a client HTTP request can travel end-to-end through the stack — gateway → one node serving all layers → valid response — before any real model or distributed logic is wired in. The six top-level packages should exist with minimal stubs: - `packages/node` — node client CLI (`meshnet-node`) - `packages/gateway` — HTTP gateway + route orchestration - `packages/tracker` — node registry and route selection - `packages/sdk` — `meshnet` Python SDK - `packages/contracts` — Solana smart contract wrappers - `packages/p2p` — gossip and shard swarm The smoke test uses a tiny in-process stub model (not a real LLM) so the test is fast and has no external dependencies. The node runs in the same process as the test. The gateway is a real HTTP server. The client sends a real `POST /v1/chat/completions` request and receives a valid OpenAI-format response.",
"acceptanceCriteria": [
"`pip install -e packages/node packages/gateway packages/tracker` works from repo root",
"`meshnet-node`, `meshnet-gateway`, and `meshnet-tracker` CLI entry points exist",
@@ -27,14 +27,14 @@
},
{
"id": "US-002",
"title": "02 \u2014 Two-node shard pipeline",
"description": "Extend the single-node smoke test so that inference is routed through exactly two nodes, each serving half of a stub model's layers. Activation tensors must travel from the gateway to node A, then from node A to node B, and the final output returns to the gateway. This proves the distributed shard pipeline works end-to-end before real models or a real tracker are introduced. The two nodes run in the same process as the test (no real network required). The gateway hardcodes the two-node route for now \u2014 dynamic route selection comes in issue 03. The activation tensor protocol (shape, dtype, serialisation format) must be established here, as all later issues depend on it. Use the domain vocabulary from `CONTEXT.md`: the ordered list of nodes is an **inference route**; each node's layer range is a **shard**; the tensors passed between nodes are activations (not \"data\" or \"chunks\").",
"title": "02 — Two-node shard pipeline",
"description": "Extend the single-node smoke test so that inference is routed through exactly two nodes, each serving half of a stub model's layers. Activation tensors must travel from the gateway to node A, then from node A to node B, and the final output returns to the gateway. This proves the distributed shard pipeline works end-to-end before real models or a real tracker are introduced. The two nodes run in the same process as the test (no real network required). The gateway hardcodes the two-node route for now — dynamic route selection comes in issue 03. The activation tensor protocol (shape, dtype, serialisation format) must be established here, as all later issues depend on it. Use the domain vocabulary from `CONTEXT.md`: the ordered list of nodes is an **inference route**; each node's layer range is a **shard**; the tensors passed between nodes are activations (not \"data\" or \"chunks\").",
"acceptanceCriteria": [
"The integration test spins up two stub nodes, each configured with a non-overlapping shard range",
"A `POST /v1/chat/completions` request results in activation tensors flowing node-A \u2192 node-B (verifiable via test assertions or logs)",
"A `POST /v1/chat/completions` request results in activation tensors flowing node-A → node-B (verifiable via test assertions or logs)",
"The gateway assembles the final response correctly from node-B's output",
"The test passes with no external services running",
"The activation tensor serialisation format is documented in a short inline comment (shape, dtype, wire format) \u2014 this becomes the contract all future nodes implement",
"The activation tensor serialisation format is documented in a short inline comment (shape, dtype, wire format) — this becomes the contract all future nodes implement",
"Follow TDD where code is added: add/adjust an observable behavior test before implementation when practical",
"Run the task-specific tests and ensure they pass",
"Run `python -m pytest` from repo root and ensure it passes, or document why no test suite exists yet",
@@ -51,8 +51,8 @@
},
{
"id": "US-003",
"title": "03 \u2014 Tracker: node registration + route selection",
"description": "Replace the hardcoded two-node route from issue 02 with a real tracker. Nodes register themselves with the tracker on startup (endpoint, shard range, hardware profile, node score). The gateway queries the tracker for an inference route for a given model preset instead of using a static list. The tracker returns an ordered list of node endpoints whose shards collectively cover all layers. The tracker runs as a lightweight HTTP service. Node score is initially a simple placeholder (e.g. fixed value or random) \u2014 real throughput/latency scoring comes later. The tracker must correctly reject route requests when no registered nodes cover a required shard range and return an appropriate error. This issue establishes the tracker's registration and routing API contract, which issues 04 (node startup) and 05 (OpenAI gateway) both depend on.",
"title": "03 — Tracker: node registration + route selection",
"description": "Replace the hardcoded two-node route from issue 02 with a real tracker. Nodes register themselves with the tracker on startup (endpoint, shard range, hardware profile, node score). The gateway queries the tracker for an inference route for a given model preset instead of using a static list. The tracker returns an ordered list of node endpoints whose shards collectively cover all layers. The tracker runs as a lightweight HTTP service. Node score is initially a simple placeholder (e.g. fixed value or random) — real throughput/latency scoring comes later. The tracker must correctly reject route requests when no registered nodes cover a required shard range and return an appropriate error. This issue establishes the tracker's registration and routing API contract, which issues 04 (node startup) and 05 (OpenAI gateway) both depend on.",
"acceptanceCriteria": [
"A node can register with the tracker via HTTP, providing its endpoint, shard range, and a hardware profile",
"The gateway queries the tracker with a model preset name and receives an ordered list of node endpoints forming a complete inference route",
@@ -76,8 +76,8 @@
},
{
"id": "US-004",
"title": "04 \u2014 Node client startup flow (`meshnet-node start`)",
"description": "Make `meshnet-node start` self-configuring from scratch on a machine with a CUDA-capable GPU (or CPU fallback). The full startup sequence must complete without any manual configuration: 1. Detect GPU model and available VRAM 2. Load an existing Solana wallet from disk, or generate and save a new one 3. Query the tracker for the optimal shard assignment given the hardware profile 4. Download the assigned shard from HuggingFace (`huggingface_hub`) 5. Register with the tracker (wallet address, endpoint, shard range, hardware profile) 6. Begin accepting inference connections The node prints a short status summary on startup: wallet address, assigned shard, model preset, and download progress. No interactive prompts \u2014 the entire flow is non-interactive. This is the primary viral growth vector. Every second of friction in this flow costs node operators. The startup sequence must work on Linux with CUDA, and degrade gracefully to CPU on machines without a GPU.",
"title": "04 — Node client startup flow (`meshnet-node start`)",
"description": "Make `meshnet-node start` self-configuring from scratch on a machine with a CUDA-capable GPU (or CPU fallback). The full startup sequence must complete without any manual configuration: 1. Detect GPU model and available VRAM 2. Load an existing Solana wallet from disk, or generate and save a new one 3. Query the tracker for the optimal shard assignment given the hardware profile 4. Download the assigned shard from HuggingFace (`huggingface_hub`) 5. Register with the tracker (wallet address, endpoint, shard range, hardware profile) 6. Begin accepting inference connections The node prints a short status summary on startup: wallet address, assigned shard, model preset, and download progress. No interactive prompts — the entire flow is non-interactive. This is the primary viral growth vector. Every second of friction in this flow costs node operators. The startup sequence must work on Linux with CUDA, and degrade gracefully to CPU on machines without a GPU.",
"acceptanceCriteria": [
"`meshnet-node start --tracker http://localhost:8080` completes startup without prompts on a machine with a CUDA GPU",
"A Solana wallet keypair is generated and saved to `~/.config/meshnet/wallet.json` if none exists",
@@ -102,8 +102,8 @@
},
{
"id": "US-005",
"title": "05 \u2014 OpenAI-compatible gateway",
"description": "Expose a production-shape HTTP API from the gateway so that any client using the OpenAI Python SDK (or any OpenAI-compatible tool) works by changing only the `base_url`. The gateway translates OpenAI-format requests into inference route execution and streams responses back in OpenAI-format. Endpoints to implement: - `POST /v1/chat/completions` \u2014 streaming (`text/event-stream`) and non-streaming - `GET /v1/models` \u2014 returns the list of model presets currently routable on the network - `GET /v1/health` \u2014 liveness check The gateway selects an inference route from the tracker, executes the shard pipeline, and assembles the streamed response. If the tracker returns no route for the requested model, the gateway responds with a standard OpenAI-format error (`model_not_available`). Authentication (API key \u2192 SOL/USDC balance) is a stub in this issue \u2014 return 200 for any non-empty `Authorization` header. Real payment gating comes in issue 06.",
"title": "05 — OpenAI-compatible gateway",
"description": "Expose a production-shape HTTP API from the gateway so that any client using the OpenAI Python SDK (or any OpenAI-compatible tool) works by changing only the `base_url`. The gateway translates OpenAI-format requests into inference route execution and streams responses back in OpenAI-format. Endpoints to implement: - `POST /v1/chat/completions` — streaming (`text/event-stream`) and non-streaming - `GET /v1/models` — returns the list of model presets currently routable on the network - `GET /v1/health` — liveness check The gateway selects an inference route from the tracker, executes the shard pipeline, and assembles the streamed response. If the tracker returns no route for the requested model, the gateway responds with a standard OpenAI-format error (`model_not_available`). Authentication (API key → SOL/USDC balance) is a stub in this issue — return 200 for any non-empty `Authorization` header. Real payment gating comes in issue 06.",
"acceptanceCriteria": [
"`openai.OpenAI(base_url=\"http://localhost:8080/v1\", api_key=\"test\").chat.completions.create(model=\"stub-model\", messages=[...])` returns a valid response",
"Streaming works: `stream=True` returns `text/event-stream` chunks in OpenAI SSE format",
@@ -127,8 +127,8 @@
},
{
"id": "US-006",
"title": "06 \u2014 Solana stake + settlement contracts",
"description": "Deploy and integrate the Solana smart contracts that make node staking, client payment, and token reward settlement trustless. All development and testing targets **Solana testnet** \u2014 never devnet or mainnet during development, to avoid real costs.",
"title": "06 — Solana stake + settlement contracts",
"description": "Deploy and integrate the Solana smart contracts that make node staking, client payment, and token reward settlement trustless. All development and testing targets **Solana testnet** — never devnet or mainnet during development, to avoid real costs.",
"acceptanceCriteria": [
"All contracts deploy successfully to Solana testnet",
"A node can submit a stake transaction and have its balance reflected in the registry contract",
@@ -136,7 +136,7 @@
"After a completed inference session, compute attribution is recorded on-chain with correct node/layer attribution",
"The epoch settlement transaction correctly distributes token rewards to node operators and deducts client balances",
"The gateway refuses to route to a node whose stake balance is below the minimum threshold",
"All contract interactions in tests run against a local Solana test validator (via `solana-test-validator`) \u2014 no live testnet required for CI",
"All contract interactions in tests run against a local Solana test validator (via `solana-test-validator`) — no live testnet required for CI",
"A `.env.testnet` config points to Solana testnet RPC for manual end-to-end testing",
"python -m pytest passes from repo root",
"Commit only this story's changes"
@@ -151,7 +151,7 @@
},
{
"id": "US-007",
"title": "07 \u2014 Fraud detection: validator + on-chain slash",
"title": "07 — Fraud detection: validator + on-chain slash",
"description": "Implement the optimistic fraud detection loop. After each inference request completes, the validator process independently decides whether to re-run the request on a reference node (~5% sample rate).",
"acceptanceCriteria": [
"The validator process samples ~5% of completed inference requests (configurable)",
@@ -175,7 +175,7 @@
},
{
"id": "US-008",
"title": "08 \u2014 Node probationary period + ban enforcement",
"title": "08 — Node probationary period + ban enforcement",
"description": "Implement the two anti-sybil mechanisms that make re-entering the network after a ban economically costly.",
"acceptanceCriteria": [
"A node with a new wallet receives no token rewards for its first N jobs (verified via settlement contract state)",
@@ -184,7 +184,7 @@
"A wallet with strike count at the threshold is marked banned in the registry contract",
"The tracker excludes banned wallets from route selection",
"A banned wallet that attempts to register with the tracker is rejected",
"An integration test covers: new wallet \u2192 N jobs \u2192 earning begins; and: strike threshold reached \u2192 banned \u2192 excluded from routes",
"An integration test covers: new wallet → N jobs → earning begins; and: strike threshold reached → banned → excluded from routes",
"python -m pytest passes from repo root",
"Commit only this story's changes"
],
@@ -198,7 +198,7 @@
},
{
"id": "US-009",
"title": "09 \u2014 P2P shard swarm",
"title": "09 — P2P shard swarm",
"description": "Once a node has downloaded a shard, it seeds that shard to other nodes that are assigned the same shard.",
"acceptanceCriteria": [
"A node that has downloaded a shard is listed as a peer for that shard in the tracker",
@@ -220,7 +220,7 @@
},
{
"id": "US-010",
"title": "10 \u2014 `meshnet` Python SDK",
"title": "10 — `meshnet` Python SDK",
"description": "A Python SDK (`packages/sdk`) that wraps the OpenAI-compatible gateway and exposes network-specific controls.",
"acceptanceCriteria": [
"`pip install meshnet` installs the SDK",
@@ -246,7 +246,7 @@
},
{
"id": "US-011",
"title": "11 \u2014 Binary activation wire format",
"title": "11 — Binary activation wire format",
"description": "Replace the base64 JSON activation payload with raw binary HTTP bodies, zstd compression, and chunked prefill (128 tokens/chunk). All nodes and the gateway must be migrated. Stub nodes continue to emit zeroed tensors, just in binary. This is a protocol prerequisite for US-012 (real model backend).",
"acceptanceCriteria": [
"Node /forward endpoint reads shape/dtype/session/chunk from HTTP headers; body is raw binary (optionally zstd-compressed)",
@@ -268,7 +268,7 @@
},
{
"id": "US-012",
"title": "12 \u2014 Real PyTorch model backend",
"title": "12 — Real PyTorch model backend",
"description": "Replace stub node inference with actual transformers layer execution. Node loads HuggingFace SafeTensors model shard (model.model.layers[start:end]), runs real forward passes in bfloat16, handles head (embed_tokens) and tail (lm_head) responsibilities, and supports bitsandbytes NF4/INT8/bfloat16 quantization via --quantization flag. Test model: openai-community/gpt2.",
"acceptanceCriteria": [
"meshnet-node start --model-id openai-community/gpt2 --shard-start 0 --shard-end 6 loads the shard without error",
@@ -283,15 +283,16 @@
"Commit only this story's changes"
],
"priority": 12,
"passes": false,
"passes": true,
"notes": "Source issue: .scratch/distributed-inference-network/issues/12-real-pytorch-model-backend.md",
"dependsOn": [
"US-011"
]
],
"completionNotes": "Completed by agent"
},
{
"id": "US-013",
"title": "13 \u2014 Coverage-first tracker shard assignment",
"title": "13 — Coverage-first tracker shard assignment",
"description": "Upgrade tracker route selection to coverage-first, speed-weighted bin-packing. Tracker maintains a live coverage map per model, issues LOAD_SHARD/DROP_SHARD rebalance directives when coverage drops, and assigns shard ranges using declared VRAM, quantization, and benchmark throughput. A model is only routable when all layer ranges have node_count >= 1.",
"acceptanceCriteria": [
"Node registration accepts vram_bytes, ram_bytes, quantizations[], benchmark_tokens_per_sec",
@@ -302,7 +303,7 @@
"Shard assignment respects VRAM: assigned_layers <= floor((vram_bytes * 0.8) / bytes_per_layer_at_quant)",
"Faster node receives wider shard range when both can cover the same gap",
"Integration test: three nodes with different VRAM show widest range on largest-VRAM node with 100% coverage",
"Integration test: kill middle-range node \u2192 tracker issues LOAD_SHARD \u2192 coverage recovers",
"Integration test: kill middle-range node → tracker issues LOAD_SHARD → coverage recovers",
"python -m pytest passes from repo root",
"Commit only this story's changes"
],
@@ -316,11 +317,11 @@
},
{
"id": "US-014",
"title": "14 \u2014 Tracker-as-first-layer-node (inference entry point)",
"title": "14 — Tracker-as-first-layer-node (inference entry point)",
"description": "Merge inference orchestration into tracker nodes that serve the first-layer shard. A tracker node exposes /v1/chat/completions directly: tokenizes input, runs embed_tokens + layers[0..k], selects onward route from coverage map, forwards binary activations, receives tail output, and streams tokens back. The standalone gateway becomes a thin load-balancer routing to tracker nodes.",
"acceptanceCriteria": [
"Node with shard_start==0 (or --tracker-mode flag) exposes /v1/chat/completions alongside /forward",
"Tracker-node /v1/chat/completions: tokenize \u2192 embed \u2192 own layers \u2192 route selection \u2192 forward binary \u2192 receive tail \u2192 stream SSE",
"Tracker-node /v1/chat/completions: tokenize → embed → own layers → route selection → forward binary → receive tail → stream SSE",
"Two tracker nodes for same model each handle requests independently",
"Gateway proxies /v1/chat/completions to tracker nodes (round-robin), discovered via GET /v1/tracker-nodes/<model_preset>",
"Integration test: two tracker nodes + two mid-shard nodes for GPT-2; 10 requests via gateway; both tracker nodes receive load; all responses valid",
@@ -338,6 +339,6 @@
}
],
"metadata": {
"updatedAt": "2026-06-29T00:00:00.000Z"
"updatedAt": "2026-06-29T12:43:53.339Z"
}
}

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

@@ -20,6 +20,18 @@ def main() -> None:
start_cmd.add_argument(
"--model", default="stub-model", help="Model preset to request from tracker"
)
start_cmd.add_argument(
"--model-id",
help="HuggingFace model id for the real PyTorch backend",
)
start_cmd.add_argument("--shard-start", type=int, help="First layer index for an explicit shard")
start_cmd.add_argument("--shard-end", type=int, help="Exclusive layer end index for an explicit shard")
start_cmd.add_argument(
"--quantization",
choices=["bfloat16", "int8", "nf4"],
default="bfloat16",
help="Weight quantization for the real PyTorch backend",
)
start_cmd.add_argument(
"--host", default="0.0.0.0", help="Interface to bind to"
)
@@ -33,13 +45,21 @@ def main() -> None:
if args.command == "start":
from meshnet_node.startup import run_startup
node = run_startup(
tracker_url=args.tracker,
port=args.port,
model=args.model,
host=args.host,
advertise_host=args.advertise_host,
)
try:
node = run_startup(
tracker_url=args.tracker,
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)

298
packages/node/meshnet_node/model_backend.py Normal file
View File

@@ -0,0 +1,298 @@
"""HuggingFace/PyTorch shard backend for real node inference."""
from __future__ import annotations
import base64
from dataclasses import dataclass
from typing import Any, Literal
Quantization = Literal["bfloat16", "int8", "nf4"]
class ModelBackendError(RuntimeError):
"""Base class for real model backend startup and execution failures."""
class MissingModelDependencyError(ModelBackendError):
"""Raised when optional model dependencies are not installed."""
class InsufficientVRAMError(ModelBackendError):
"""Raised when a requested shard cannot fit in available CUDA memory."""
@dataclass(frozen=True)
class TensorPayload:
body: bytes
shape: list[int]
attention_mask_header: str | None
position_ids_header: str | None
def validate_quantization(value: str) -> Quantization:
if value not in {"bfloat16", "int8", "nf4"}:
raise ValueError("quantization must be one of: bfloat16, int8, nf4")
return value # type: ignore[return-value]
def build_quantization_config(quantization: Quantization) -> Any | None:
"""Return a transformers BitsAndBytesConfig for quantized weights."""
if quantization == "bfloat16":
return None
try:
import torch
from transformers import BitsAndBytesConfig
except ModuleNotFoundError as exc:
raise MissingModelDependencyError(
"transformers and torch are required for int8/nf4 quantization"
) from exc
if quantization == "int8":
return BitsAndBytesConfig(load_in_8bit=True)
return BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
class TorchModelShard:
"""Executable subset of a HuggingFace causal language model."""
def __init__(
self,
model_id: str,
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
) -> None:
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
raise ValueError("shard_start must be <= shard_end and non-negative")
self.model_id = model_id
self.shard_start = shard_start
self.shard_end = shard_end
self.quantization = quantization
try:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
except ModuleNotFoundError as exc:
raise MissingModelDependencyError(
"real model backend requires torch, transformers, safetensors, accelerate, and bitsandbytes"
) from exc
self.torch = torch
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
quant_config = build_quantization_config(quantization)
try:
self.model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=quant_config,
device_map="auto" if quant_config is not None else None,
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True,
use_safetensors=True,
)
if quant_config is None:
self.model.to(self.device)
except Exception as exc:
if _looks_like_oom(exc):
raise InsufficientVRAMError(
f"insufficient VRAM to load {model_id} layers {shard_start}:{shard_end} "
f"with {quantization} quantization; choose a smaller shard or lower quantization"
) from exc
raise
self.model.eval()
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
self.layers = _model_layers(self.model)
self.total_layers = len(self.layers)
if shard_end > self.total_layers:
raise ValueError(
f"shard_end {shard_end} exceeds total layer count {self.total_layers}"
)
self.is_head = shard_start == 0
self.is_tail = shard_end == self.total_layers
self.hidden_size = int(
getattr(self.model.config, "hidden_size", 0)
or getattr(self.model.config, "n_embd", 0)
)
self._embed_tokens = _embed_tokens(self.model) if self.is_head else None
self._position_embeddings = _position_embeddings(self.model)
self._norm = _final_norm(self.model) if self.is_tail else None
self._lm_head = getattr(self.model, "lm_head", None) if self.is_tail else None
def encode_prompt(self, prompt: str) -> TensorPayload:
if not self.is_head or self._embed_tokens is None:
raise ModelBackendError("text prompts can only be accepted by the head shard")
encoded = self.tokenizer(prompt, return_tensors="pt")
input_ids = encoded["input_ids"].to(self.device)
attention_mask = encoded.get("attention_mask")
if attention_mask is None:
attention_mask = self.torch.ones_like(input_ids)
attention_mask = attention_mask.to(self.device)
position_ids = _position_ids(attention_mask, self.torch)
hidden_states = self._embed_tokens(input_ids)
if self._position_embeddings is not None:
hidden_states = hidden_states + self._position_embeddings(position_ids)
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
return self._payload(hidden_states, attention_mask, position_ids)
def forward_bytes(
self,
body: bytes,
shape: list[int],
attention_mask_header: str | None,
position_ids_header: str | None,
) -> TensorPayload | str:
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
self.device
)
attention_mask = _tensor_from_int64_header(
attention_mask_header, self.torch, self.device
)
position_ids = _tensor_from_int64_header(
position_ids_header, self.torch, self.device
)
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
if self.is_tail:
return self.decode_tail(hidden_states)
return self._payload(hidden_states, attention_mask, position_ids)
def decode_tail(self, hidden_states: Any) -> str:
if self._norm is not None:
hidden_states = self._norm(hidden_states)
if self._lm_head is None:
raise ModelBackendError("tail shard has no lm_head")
logits = self._lm_head(hidden_states)
token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
return self.tokenizer.decode([token_id], skip_special_tokens=True)
def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
with self.torch.inference_mode():
for layer in self.layers[self.shard_start:self.shard_end]:
hidden_states = _call_layer(layer, hidden_states, attention_mask, position_ids)
return hidden_states.to(self.torch.bfloat16)
def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload:
hidden_states = hidden_states.to(self.torch.bfloat16).contiguous()
return TensorPayload(
body=_tensor_to_bytes(hidden_states),
shape=list(hidden_states.shape),
attention_mask_header=_int_tensor_header(attention_mask)
if attention_mask is not None
else None,
position_ids_header=_int_tensor_header(position_ids)
if position_ids is not None
else None,
)
def load_torch_shard(
model_id: str,
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
) -> TorchModelShard:
return TorchModelShard(model_id, shard_start, shard_end, quantization)
def _model_layers(model: Any) -> Any:
if hasattr(model, "model") and hasattr(model.model, "layers"):
return model.model.layers
if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
return model.transformer.h
raise ModelBackendError(
"unsupported HuggingFace model architecture: no transformer layers found"
)
def _embed_tokens(model: Any) -> Any:
if hasattr(model, "model") and hasattr(model.model, "embed_tokens"):
return model.model.embed_tokens
if hasattr(model, "transformer") and hasattr(model.transformer, "wte"):
return model.transformer.wte
raise ModelBackendError(
"unsupported HuggingFace model architecture: no token embeddings found"
)
def _position_embeddings(model: Any) -> Any | None:
if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
return model.transformer.wpe
return None
def _final_norm(model: Any) -> Any | None:
if hasattr(model, "model") and hasattr(model.model, "norm"):
return model.model.norm
if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"):
return model.transformer.ln_f
return None
def _position_ids(attention_mask: Any, torch: Any) -> Any:
position_ids = attention_mask.long().cumsum(-1) - 1
return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long)
def _call_layer(layer: Any, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
attempts = (
{
"attention_mask": attention_mask,
"position_ids": position_ids,
"use_cache": False,
},
{"attention_mask": attention_mask, "use_cache": False},
{"use_cache": False},
{},
)
last_exc: Exception | None = None
for kwargs in attempts:
filtered = {key: value for key, value in kwargs.items() if value is not None}
try:
output = layer(hidden_states, **filtered)
return output[0] if isinstance(output, tuple) else output
except TypeError as exc:
last_exc = exc
if last_exc is not None:
raise last_exc
return layer(hidden_states)[0]
def _tensor_to_bytes(tensor: Any) -> bytes:
import torch
return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()
def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any:
tensor = torch.frombuffer(bytearray(body), dtype=torch.bfloat16)
return tensor.reshape(shape)
def _int_tensor_header(tensor: Any) -> str:
data = tensor.detach().cpu().to(tensor.int64).contiguous()
raw = data.numpy().tobytes()
shape = ",".join(str(dim) for dim in data.shape)
encoded = base64.b64encode(raw).decode("ascii")
return f"{shape}:{encoded}"
def _tensor_from_int64_header(value: str | None, torch: Any, device: Any) -> Any | None:
if not value:
return None
shape_text, encoded = value.split(":", 1)
shape = [int(part) for part in shape_text.split(",") if part]
raw = base64.b64decode(encoded.encode("ascii"))
return torch.frombuffer(bytearray(raw), dtype=torch.int64).reshape(shape).to(device)
def _looks_like_oom(exc: BaseException) -> bool:
current: BaseException | None = exc
while current is not None:
text = str(current).lower()
if "out of memory" in text or "cuda error: out of memory" in text:
return True
current = current.__cause__ or current.__context__
return False

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

@@ -14,6 +14,7 @@ from typing import Any
from .downloader import compute_shard_checksum, download_shard
from .hardware import detect_hardware
from .server import StubNodeServer
from .torch_server import TorchNodeServer
from .wallet import load_or_create_wallet
@@ -35,12 +36,16 @@ def run_startup(
tracker_url: str,
port: int = 0,
model: str = "stub-model",
model_id: str | None = None,
shard_start: int | None = None,
shard_end: int | None = None,
quantization: str = "bfloat16",
wallet_path: Path | None = None,
cache_dir: Path | None = None,
host: str = "127.0.0.1",
advertise_host: str | None = None,
contracts: Any | None = None,
) -> StubNodeServer:
) -> StubNodeServer | TorchNodeServer:
"""Execute the full startup sequence and return a running node server.
Steps (all non-interactive):
@@ -79,6 +84,35 @@ def run_startup(
if probationary_line is not None:
print(f" {probationary_line}", flush=True)
if model_id is not None and shard_start is not None and shard_end is not None:
print("Loading real PyTorch model shard...", flush=True)
node = TorchNodeServer(
host=host,
port=port,
model_id=model_id,
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
)
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}"
print(
f"\n{'=' * 32}\n"
f"meshnet-node ready\n"
f" Wallet: {address}\n"
f" Model ID: {model_id}\n"
f" Shard: layers {shard_start}-{shard_end}\n"
f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n"
f" Hardware: {device.upper()}\n"
f"{'=' * 32}",
flush=True,
)
return node
if model_id is not None or shard_start is not None or shard_end is not None:
raise ValueError("--model-id, --shard-start, and --shard-end must be provided together")
# 3. Shard assignment from tracker
print("Querying tracker for shard assignment...", flush=True)
assign_qs = urllib.parse.urlencode({

271
packages/node/meshnet_node/torch_server.py Normal file
View File

@@ -0,0 +1,271 @@
"""HTTP server for real PyTorch-backed shard nodes."""
from __future__ import annotations
import http.server
import json
import sys
import threading
from .model_backend import (
InsufficientVRAMError,
MissingModelDependencyError,
Quantization,
TorchModelShard,
validate_quantization,
)
from .server import (
_WIRE_VERSION,
_compress_body,
_decompress_body,
_parse_shape,
_validate_activation_body,
)
class _TorchHTTPServer(http.server.HTTPServer):
def __init__(self, addr, handler, backend: TorchModelShard):
super().__init__(addr, handler)
self.backend = backend
self.received_activations = False
self.forward_chunk_count = 0
class _TorchHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
pass
def do_POST(self):
if self.path == "/forward":
self._handle_forward()
elif self.path == "/v1/infer":
self._handle_infer()
else:
self.send_response(404)
self.end_headers()
def _handle_infer(self) -> None:
body = self._read_json_body()
if body is None:
return
messages = body.get("messages", [])
prompt = ""
if isinstance(messages, list) and messages:
last = messages[-1]
if isinstance(last, dict):
prompt = str(last.get("content", ""))
server: _TorchHTTPServer = self.server # type: ignore[assignment]
try:
payload = server.backend.encode_prompt(prompt)
if server.backend.is_tail:
text = server.backend.decode_tail(
server.backend.torch.frombuffer(
bytearray(payload.body),
dtype=server.backend.torch.bfloat16,
)
.reshape(payload.shape)
.to(server.backend.device)
)
self._send_json(200, {"text": text})
return
self._send_json(200, {"activations": {"shape": payload.shape, "dtype": "bfloat16"}})
except Exception as exc:
self._send_json(500, {"error": str(exc)})
def _handle_forward(self) -> None:
content_type = self.headers.get("Content-Type", "")
if content_type.startswith("application/json"):
self._handle_prompt_forward()
return
self._handle_binary_forward()
def _handle_prompt_forward(self) -> None:
body = self._read_json_body()
if body is None:
return
prompt = str(body.get("prompt", ""))
server: _TorchHTTPServer = self.server # type: ignore[assignment]
try:
payload = server.backend.encode_prompt(prompt)
except Exception as exc:
self._send_json(400, {"error": str(exc)})
return
self._send_activation(payload)
def _handle_binary_forward(self) -> None:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
try:
shape = _parse_shape(self.headers.get("X-Meshnet-Shape"))
dtype = self.headers.get("X-Meshnet-Dtype", "")
session = self.headers["X-Meshnet-Session"]
chunk_index = self.headers["X-Meshnet-Chunk-Index"]
chunk_total = self.headers["X-Meshnet-Chunk-Total"]
encoding = self.headers.get("X-Meshnet-Encoding")
length = int(self.headers.get("Content-Length", 0))
body = self.rfile.read(length)
raw_body = _decompress_body(body, encoding)
_validate_activation_body(raw_body, shape, dtype)
if dtype != "bfloat16":
raise ValueError("real model backend requires bfloat16 activation input")
chunk_index_value = int(chunk_index)
chunk_total_value = int(chunk_total)
if chunk_total_value <= 0 or not 0 <= chunk_index_value < chunk_total_value:
raise ValueError("invalid chunk index/total")
except (KeyError, ValueError, TypeError):
self.send_response(400)
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
self.end_headers()
return
server.forward_chunk_count += 1
if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
server.received_activations = True
try:
result = server.backend.forward_bytes(
raw_body,
shape,
self.headers.get("X-Meshnet-Attn-Mask"),
self.headers.get("X-Meshnet-Position-Ids"),
)
except Exception as exc:
self._send_json(500, {"error": str(exc)})
return
if isinstance(result, str):
self._send_json(200, {"text": result})
return
response_body = _compress_body(result.body, encoding)
self.send_response(200)
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(len(response_body)))
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in result.shape))
self.send_header("X-Meshnet-Dtype", "bfloat16")
self.send_header("X-Meshnet-Session", session)
self.send_header("X-Meshnet-Chunk-Index", chunk_index)
self.send_header("X-Meshnet-Chunk-Total", chunk_total)
if encoding:
self.send_header("X-Meshnet-Encoding", encoding)
if result.attention_mask_header:
self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header)
if result.position_ids_header:
self.send_header("X-Meshnet-Position-Ids", result.position_ids_header)
self.end_headers()
self.wfile.write(response_body)
def _send_activation(self, payload) -> None:
body = payload.body
self.send_response(200)
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(len(body)))
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in payload.shape))
self.send_header("X-Meshnet-Dtype", "bfloat16")
if payload.attention_mask_header:
self.send_header("X-Meshnet-Attn-Mask", payload.attention_mask_header)
if payload.position_ids_header:
self.send_header("X-Meshnet-Position-Ids", payload.position_ids_header)
self.end_headers()
self.wfile.write(body)
def _read_json_body(self) -> dict | None:
length = int(self.headers.get("Content-Length", 0))
try:
body = json.loads(self.rfile.read(length) or b"{}")
except (json.JSONDecodeError, ValueError):
self._send_json(400, {"error": "invalid JSON body"})
return None
if not isinstance(body, dict):
self._send_json(400, {"error": "JSON body must be an object"})
return None
return body
def _send_json(self, status: int, data: dict) -> None:
payload = json.dumps(data).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
class TorchNodeServer:
"""HTTP server backed by a HuggingFace causal language model shard."""
def __init__(
self,
host: str = "127.0.0.1",
port: int = 0,
model_id: str = "openai-community/gpt2",
shard_start: int = 0,
shard_end: int = 6,
quantization: str = "bfloat16",
backend: TorchModelShard | None = None,
) -> None:
self._host = host
self._requested_port = port
self._backend = backend or _load_backend(
model_id,
shard_start,
shard_end,
quantization,
)
self._server: _TorchHTTPServer | None = None
self._thread: threading.Thread | None = None
self.port: int | None = None
@property
def backend(self) -> TorchModelShard:
return self._backend
@property
def received_activations(self) -> bool:
return self._server.received_activations if self._server is not None else False
@property
def forward_chunk_count(self) -> int:
return self._server.forward_chunk_count if self._server is not None else 0
def start(self) -> int:
if self._server is not None:
raise RuntimeError("TorchNodeServer is already running")
self._server = _TorchHTTPServer(
(self._host, self._requested_port),
_TorchHandler,
self._backend,
)
self.port = self._server.server_address[1]
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
self._thread.start()
return self.port
def stop(self) -> None:
if self._server is None:
return
self._server.shutdown()
self._server.server_close()
if self._thread is not None:
self._thread.join(timeout=1)
self._server = None
self._thread = None
self.port = None
def _load_backend(
model_id: str,
shard_start: int,
shard_end: int,
quantization: str,
) -> TorchModelShard:
from .model_backend import load_torch_shard
quant = validate_quantization(quantization)
try:
return load_torch_shard(model_id, shard_start, shard_end, quant)
except MissingModelDependencyError:
raise
except InsufficientVRAMError as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
raise

5
packages/node/pyproject.toml
View File

@@ -11,6 +11,11 @@ requires-python = ">=3.10"
dependencies = [
"cryptography>=41",
"huggingface-hub>=0.20",
"accelerate>=0.28",
"bitsandbytes>=0.43",
"safetensors>=0.4",
"torch>=2.1",
"transformers>=4.39",
"zstandard>=0.22",
]

3
pyproject.toml
View File

@@ -13,3 +13,6 @@ dev = ["pytest>=8", "openai>=1", "langchain-openai>=0.1"]
[tool.pytest.ini_options]
testpaths = ["tests"]
markers = [
"integration: tests that download models, require GPUs, or exercise external integrations",
]

208
tests/test_real_model_backend.py Normal file
View File

@@ -0,0 +1,208 @@
"""US-012 tests for the real PyTorch node backend."""
import json
import os
from pathlib import Path
import sys
import types
import urllib.request
import pytest
from meshnet_node.model_backend import (
InsufficientVRAMError,
TensorPayload,
build_quantization_config,
validate_quantization,
)
from meshnet_node.torch_server import TorchNodeServer
class _FakeBackend:
model_id = "fake-model"
total_layers = 12
is_head = True
is_tail = False
def encode_prompt(self, prompt: str) -> TensorPayload:
assert prompt == "The capital of France is"
return TensorPayload(
body=b"\x00" * (1 * 6 * 8 * 2),
shape=[1, 6, 8],
attention_mask_header=None,
position_ids_header=None,
)
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header):
assert shape == [1, 6, 8]
return TensorPayload(
body=body,
shape=shape,
attention_mask_header=attention_mask_header,
position_ids_header=position_ids_header,
)
class _FakeTailBackend(_FakeBackend):
is_head = False
is_tail = True
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header):
assert len(body) == 1 * 6 * 8 * 2
return " Paris"
def test_quantization_flag_validation():
assert validate_quantization("bfloat16") == "bfloat16"
assert validate_quantization("int8") == "int8"
assert validate_quantization("nf4") == "nf4"
with pytest.raises(ValueError, match="quantization"):
validate_quantization("float32")
def test_node_package_declares_torch_dependency():
pyproject = Path("packages/node/pyproject.toml").read_text(encoding="utf-8")
assert '"torch>=' in pyproject
def test_bitsandbytes_configs_are_created_lazily(monkeypatch):
calls = []
class FakeBitsAndBytesConfig:
def __init__(self, **kwargs):
calls.append(kwargs)
monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(bfloat16="bf16"))
monkeypatch.setitem(
sys.modules,
"transformers",
types.SimpleNamespace(BitsAndBytesConfig=FakeBitsAndBytesConfig),
)
assert build_quantization_config("bfloat16") is None
build_quantization_config("int8")
build_quantization_config("nf4")
assert calls == [
{"load_in_8bit": True},
{
"load_in_4bit": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": "bf16",
},
]
def test_head_forward_accepts_text_prompt_and_returns_bfloat16_activations():
node = TorchNodeServer(backend=_FakeBackend())
port = node.start()
try:
payload = json.dumps({"prompt": "The capital of France is"}).encode()
req = urllib.request.Request(
f"http://127.0.0.1:{port}/forward",
data=payload,
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req, timeout=5) as resp:
body = resp.read()
headers = {key.lower(): value for key, value in resp.headers.items()}
assert len(body) == 1 * 6 * 8 * 2
assert headers["x-meshnet-shape"] == "1,6,8"
assert headers["x-meshnet-dtype"] == "bfloat16"
assert headers["x-meshnet-wire"] == "2"
finally:
node.stop()
def test_tail_forward_returns_text_completion_from_binary_activations():
node = TorchNodeServer(backend=_FakeTailBackend())
port = node.start()
try:
req = urllib.request.Request(
f"http://127.0.0.1:{port}/forward",
data=b"\x00" * (1 * 6 * 8 * 2),
headers={
"Content-Type": "application/octet-stream",
"X-Meshnet-Shape": "1,6,8",
"X-Meshnet-Dtype": "bfloat16",
"X-Meshnet-Session": "session-1",
"X-Meshnet-Chunk-Index": "0",
"X-Meshnet-Chunk-Total": "1",
"X-Meshnet-Hop-Index": "1",
},
method="POST",
)
with urllib.request.urlopen(req, timeout=5) as resp:
body = json.loads(resp.read())
assert body == {"text": " Paris"}
assert node.received_activations
assert node.forward_chunk_count == 1
finally:
node.stop()
@pytest.mark.integration
def test_two_node_gpt2_completion_is_deterministic():
if os.environ.get("CI"):
pytest.skip("GPT-2 integration test is skipped in CI")
torch = pytest.importorskip("torch")
pytest.importorskip("transformers")
pytest.importorskip("safetensors")
pytest.importorskip("accelerate")
pytest.importorskip("bitsandbytes")
if not torch.cuda.is_available():
pytest.skip("GPT-2 integration test requires a CUDA GPU")
head = TorchNodeServer(
model_id="openai-community/gpt2",
shard_start=0,
shard_end=6,
quantization="bfloat16",
)
tail = TorchNodeServer(
model_id="openai-community/gpt2",
shard_start=6,
shard_end=12,
quantization="bfloat16",
)
head_port = head.start()
tail_port = tail.start()
try:
prompt_req = urllib.request.Request(
f"http://127.0.0.1:{head_port}/forward",
data=json.dumps({"prompt": "The capital of France is"}).encode(),
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(prompt_req, timeout=60) as resp:
activation = resp.read()
head_headers = resp.headers
tail_req = urllib.request.Request(
f"http://127.0.0.1:{tail_port}/forward",
data=activation,
headers={
"Content-Type": "application/octet-stream",
"X-Meshnet-Shape": head_headers["X-Meshnet-Shape"],
"X-Meshnet-Dtype": head_headers["X-Meshnet-Dtype"],
"X-Meshnet-Session": "gpt2-session",
"X-Meshnet-Chunk-Index": "0",
"X-Meshnet-Chunk-Total": "1",
"X-Meshnet-Hop-Index": "1",
"X-Meshnet-Attn-Mask": head_headers["X-Meshnet-Attn-Mask"],
"X-Meshnet-Position-Ids": head_headers["X-Meshnet-Position-Ids"],
},
method="POST",
)
with urllib.request.urlopen(tail_req, timeout=60) as resp:
body = json.loads(resp.read())
assert body["text"].strip()
assert body["text"] == " Paris"
finally:
head.stop()
tail.stop()