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

feat: checkpoint distributed gguf runtime stories

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-15 23:42:58 +03:00
parent eaf00f6add
commit 1fe31ef38d
60 changed files with 8478 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -29,6 +29,7 @@ from .model_catalog import model_metadata_for
from .recipe_manifest import DEFAULT_RECIPE_ID, Recipe, RecipeManifest, load_recipe_manifest
from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
from .server import StubNodeServer
from .gguf_backend import build_gguf_backend
from .torch_server import TorchNodeServer
from .wallet import load_or_create_wallet
@@ -662,6 +663,35 @@ def _resolve_recipe(recipe_id: str | None) -> tuple[RecipeManifest, Recipe]:
return manifest, manifest.require(recipe_id or DEFAULT_RECIPE_ID)
def _gguf_backend_for_recipe(
recipe: Recipe,
*,
model_id: str,
shard_start: int,
shard_end: int,
quantization: str,
total_layers: int | None,
device: str,
model_revision: str | None = None,
) -> object | None:
"""Build the GGUF backend only for recipes that explicitly ask for it."""
if recipe.backend_id != "llama.cpp":
return None
return build_gguf_backend(
model_id=model_id,
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
total_layers=total_layers,
model_revision=model_revision,
device_type=device,
architecture_adapter="dense-llama",
tokenizer_revision=model_revision or model_id,
runtime_recipe_fingerprint=None,
supports_kv_cache=recipe.params.get("use_cache", True) is not False,
)
def _capability_device(backend: Any, detected_device: str) -> str:
"""The device the shard actually landed on, or the one this node detected."""
device = getattr(backend, "device", None)
@@ -875,7 +905,8 @@ def run_startup(
if model_id: # treat "" the same as None — no explicit model given
full_sources: list[dict] = []
# Auto-detect shard range from model config if not explicitly provided
detected: int | None = None
# Auto-detect shard range from model config if not explicitly provided.
if shard_start is None or shard_end is None:
try:
detected = _detect_num_layers(model_id, cache_dir=cache_dir)
@@ -939,22 +970,38 @@ def run_startup(
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)
node = TorchNodeServer(
host=host,
port=port,
backend = _gguf_backend_for_recipe(
recipe,
model_id=model_id,
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
tracker_url=tracker_url,
route_timeout=route_timeout,
cache_dir=cache_dir,
debug=debug,
max_loaded_shards=max_loaded_shards,
force_cpu=force_cpu,
recipe_params=recipe.params,
total_layers=detected if detected is not None else (shard_end + 1 if shard_end is not None else None),
device=device,
model_revision=None,
)
print(
"Loading native llama.cpp model shard..." if backend is not None else "Loading real PyTorch model shard...",
flush=True,
)
node_kwargs = {
"host": host,
"port": port,
"model_id": model_id,
"shard_start": shard_start,
"shard_end": shard_end,
"quantization": quantization,
"tracker_url": tracker_url,
"route_timeout": route_timeout,
"cache_dir": cache_dir,
"debug": debug,
"max_loaded_shards": max_loaded_shards,
"force_cpu": force_cpu,
"recipe_params": recipe.params,
}
if backend is not None:
node_kwargs["backend"] = backend
node = TorchNodeServer(**node_kwargs)
capability_report = _admit_capability(
node,
model_id=model_id,
@@ -968,10 +1015,15 @@ def run_startup(
recipe=recipe,
validator=capability_validator,
)
proof_shard = capability_report.shard
_node_start_time = time.monotonic()
actual_port = node.start()
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
shard_label = _format_shard_label(shard_start, shard_end, total_layers)
shard_label = _format_shard_label(
proof_shard.start,
proof_shard.end,
total_layers,
)
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}"
if hasattr(node, "set_advertised_endpoint"):
@@ -994,16 +1046,17 @@ def run_startup(
"model": model_id.split("/")[-1],
"hf_repo": model_id,
"num_layers": total_layers,
"shard_start": shard_start,
"shard_end": shard_end,
"shard_start": proof_shard.start,
"shard_end": proof_shard.end,
"hardware_profile": hw,
"wallet_address": address,
"quantization": quantization,
"score": 1.0,
"tracker_mode": (shard_start == 0),
"tracker_mode": (proof_shard.start == 0),
"managed_assignment": not user_pinned_shard,
"model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
"capability_report": capability_report.to_dict(),
"compatibility_fingerprint": capability_report.compatibility_fingerprint,
# Declared independently of the proof: the tracker checks that the
# recipe this node says it serves with is the one the proof ran.
"recipe_id": recipe.id,
@@ -1011,8 +1064,8 @@ def run_startup(
"downloaded_models": (
_downloaded_model_inventory(
model_id.split("/")[-1],
shard_start,
shard_end,
proof_shard.start,
proof_shard.end,
model_cache_path,
hf_repo=model_id,
model_sources=full_sources,
@@ -1114,22 +1167,38 @@ def run_startup(
hf_repo=assigned_hf_repo,
model_sources=full_sources,
)
print("Loading real PyTorch model shard...", flush=True)
node = TorchNodeServer(
host=host,
port=port,
backend = _gguf_backend_for_recipe(
recipe,
model_id=assigned_hf_repo,
shard_start=assigned_shard_start,
shard_end=assigned_shard_end,
quantization=quantization,
tracker_url=tracker_url,
route_timeout=route_timeout,
cache_dir=cache_dir,
debug=debug,
max_loaded_shards=max_loaded_shards,
force_cpu=force_cpu,
recipe_params=recipe.params,
total_layers=assigned_num_layers,
device=device,
model_revision=None,
)
print(
"Loading native llama.cpp model shard..." if backend is not None else "Loading real PyTorch model shard...",
flush=True,
)
node_kwargs = {
"host": host,
"port": port,
"model_id": assigned_hf_repo,
"shard_start": assigned_shard_start,
"shard_end": assigned_shard_end,
"quantization": quantization,
"tracker_url": tracker_url,
"route_timeout": route_timeout,
"cache_dir": cache_dir,
"debug": debug,
"max_loaded_shards": max_loaded_shards,
"force_cpu": force_cpu,
"recipe_params": recipe.params,
}
if backend is not None:
node_kwargs["backend"] = backend
node = TorchNodeServer(**node_kwargs)
capability_report = _admit_capability(
node,
model_id=assigned_hf_repo,
@@ -1143,6 +1212,7 @@ def run_startup(
recipe=recipe,
validator=capability_validator,
)
proof_shard = capability_report.shard
_node_start_time = time.monotonic()
actual_port = node.start()
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
@@ -1165,16 +1235,17 @@ def run_startup(
"model": assigned_hf_repo.split("/")[-1],
"hf_repo": assigned_hf_repo,
"num_layers": assigned_num_layers,
"shard_start": assigned_shard_start,
"shard_end": assigned_shard_end,
"shard_start": proof_shard.start,
"shard_end": proof_shard.end,
"hardware_profile": hw,
"wallet_address": address,
"quantization": quantization,
"score": 1.0,
"tracker_mode": (assigned_shard_start == 0),
"tracker_mode": (proof_shard.start == 0),
"managed_assignment": True,
"model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
"capability_report": capability_report.to_dict(),
"compatibility_fingerprint": capability_report.compatibility_fingerprint,
# Declared independently of the proof: the tracker checks that the
# recipe this node says it serves with is the one the proof ran.
"recipe_id": recipe.id,
@@ -1182,8 +1253,8 @@ def run_startup(
"downloaded_models": (
_downloaded_model_inventory(
assigned_hf_repo.split("/")[-1],
assigned_shard_start,
assigned_shard_end,
proof_shard.start,
proof_shard.end,
model_cache_path,
hf_repo=assigned_hf_repo,
model_sources=full_sources,
@@ -1199,8 +1270,8 @@ def run_startup(
tracker_url, auto_reg_payload, node, _node_start_time,
)
shard_label = _format_shard_label(
assigned_shard_start,
assigned_shard_end,
proof_shard.start,
proof_shard.end,
assigned_num_layers,
)
print(
@@ -1315,22 +1386,38 @@ def run_startup(
# 5. Start HTTP server — real HF weights use TorchNodeServer; stub-model stays stub.
_node_start_time = time.monotonic()
if hf_repo and assigned_model != "stub-model":
print("Loading real PyTorch model shard...", flush=True)
node = TorchNodeServer(
host=host,
port=port,
backend = _gguf_backend_for_recipe(
recipe,
model_id=hf_repo,
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
tracker_url=tracker_url,
route_timeout=route_timeout,
cache_dir=shard_path,
debug=debug,
max_loaded_shards=max_loaded_shards,
force_cpu=force_cpu,
recipe_params=recipe.params,
total_layers=total_layers,
device=device,
model_revision=None,
)
print(
"Loading native llama.cpp model shard..." if backend is not None else "Loading real PyTorch model shard...",
flush=True,
)
node_kwargs = {
"host": host,
"port": port,
"model_id": hf_repo,
"shard_start": shard_start,
"shard_end": shard_end,
"quantization": quantization,
"tracker_url": tracker_url,
"route_timeout": route_timeout,
"cache_dir": shard_path,
"debug": debug,
"max_loaded_shards": max_loaded_shards,
"force_cpu": force_cpu,
"recipe_params": recipe.params,
}
if backend is not None:
node_kwargs["backend"] = backend
node = TorchNodeServer(**node_kwargs)
capability_report = _admit_capability(
node,
model_id=hf_repo,
@@ -1379,6 +1466,7 @@ def run_startup(
"managed_assignment": not user_pinned_shard,
"model_metadata": model_metadata_for(hf_repo, total_layers, cache_dir=shard_path),
"capability_report": capability_report.to_dict(),
"compatibility_fingerprint": capability_report.compatibility_fingerprint,
# Declared independently of the proof: the tracker checks that the
# recipe this node says it serves with is the one the proof ran.
"recipe_id": recipe.id,
@@ -1431,6 +1519,7 @@ def run_startup(
recipe=recipe,
validator=capability_validator,
)
proof_shard = capability_report.shard
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}"
@@ -1450,10 +1539,11 @@ def run_startup(
reg_payload = {
"endpoint": endpoint,
"model": assigned_model,
"shard_start": shard_start,
"shard_end": shard_end,
"shard_start": proof_shard.start,
"shard_end": proof_shard.end,
"shard_checksum": shard_checksum,
"capability_report": capability_report.to_dict(),
"compatibility_fingerprint": capability_report.compatibility_fingerprint,
# Declared independently of the proof: the tracker checks that the
# recipe this node says it serves with is the one the proof ran.
"recipe_id": recipe.id,
@@ -1484,8 +1574,8 @@ def run_startup(
if gpu_name:
hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
shard_label = _format_shard_label(
shard_start,
shard_end,
proof_shard.start,
proof_shard.end,
assigned_total_layers,
model_name=assigned_model,
)