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

fix model registration and anouncement. added console panel

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-07 14:24:37 +02:00
parent 0e8acf5d59
commit 339577a26c
9 changed files with 732 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
packages/node/meshnet_node/downloader.py
View File

@@ -1,6 +1,6 @@
"""Shard downloader — fetches model shards from peers or HuggingFace Hub.
Cache layout: ~/.cache/meshnet/shards/<model>/layers_<start>-<end>/
Cache layout: ~/.cache/meshnet/shards/<model>/
For "stub-model" (no HF repo), a placeholder JSON file is written so the
test suite never touches the network.
@@ -106,9 +106,7 @@ def _download_shard_from_peer(
_safe_extract_shard(archive_path, extract_dir)
if compute_shard_checksum(extract_dir) != checksum:
return False
if shard_dir.exists():
shutil.rmtree(shard_dir)
shutil.move(str(extract_dir), str(shard_dir))
_merge_tree(extract_dir, shard_dir)
return True
except Exception:
return False
@@ -217,6 +215,38 @@ def _reuse_local_file(expected: int, dest: Path, final: Path) -> bool:
return False
def _valid_source_rel_files(source: dict) -> list[str]:
return [
rel for rel in (source.get("files") or [])
if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts
]
def _source_files_cached(source: dict, shard_dir: Path) -> bool:
rel_files = _valid_source_rel_files(source)
if not rel_files:
return False
sizes = source.get("file_sizes")
if not isinstance(sizes, dict):
return False
for rel in rel_files:
expected = sizes.get(rel)
path = shard_dir / rel
if not isinstance(expected, int) or not path.exists() or path.stat().st_size != expected:
return False
return True
def _merge_tree(src: Path, dest: Path) -> None:
dest.mkdir(parents=True, exist_ok=True)
for path in sorted(p for p in src.rglob("*") if p.is_file()):
target = dest / path.relative_to(src)
target.parent.mkdir(parents=True, exist_ok=True)
if target.exists():
target.unlink()
shutil.move(str(path), str(target))
def _download_source_files(
source: dict,
shard_dir: Path,
@@ -231,10 +261,7 @@ def _download_source_files(
the download is retried or the node is restarted.
"""
url = source.get("url")
rel_files = [
rel for rel in (source.get("files") or [])
if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts
]
rel_files = _valid_source_rel_files(source)
if not isinstance(url, str) or not url or not rel_files:
return None
sizes = source.get("file_sizes")
@@ -283,9 +310,11 @@ def _download_source_files(
time.sleep(1.0 * attempt)
finally:
tracker_bar.close()
if shard_dir.exists():
shutil.rmtree(shard_dir)
shutil.move(str(partial_dir), str(shard_dir))
_merge_tree(partial_dir, shard_dir)
try:
partial_dir.rmdir()
except OSError:
pass
return shard_dir
@@ -331,9 +360,7 @@ def _download_model_source(
if progress:
print(f" {label}: transfer complete ({received / 1e9:.2f} GB), extracting ...", flush=True)
_safe_extract_shard(archive_path, extract_dir)
if shard_dir.exists():
shutil.rmtree(shard_dir)
shutil.move(str(extract_dir), str(shard_dir))
_merge_tree(extract_dir, shard_dir)
return shard_dir
except Exception as exc:
if progress:
@@ -431,10 +458,17 @@ def download_shard(
the test suite hermetic while the real download path is exercised by
passing a non-stub *hf_repo*.
"""
shard_dir = cache_dir / model / f"layers_{shard_start}-{shard_end}"
shard_dir = cache_dir / model
if progress:
print(f" Target location: {shard_dir}", flush=True)
for source in model_sources or []:
label = str(source.get("type") or "model-source")
if _source_files_cached(source, shard_dir):
if progress:
print(f" [{label}] requested files already cached at {shard_dir}", flush=True)
return shard_dir
for peer in peers or []:
if progress:
print(f" Trying peer shard download from {peer.get('endpoint')} ...", flush=True)

17
packages/node/meshnet_node/model_catalog.py
View File

@@ -64,6 +64,17 @@ def _load_model_metadata() -> dict[str, dict]:
_MODEL_METADATA = _load_model_metadata()
def _local_model_path(hf_repo: str, cache_dir: Path | None) -> Path | None:
if cache_dir is None:
return None
if (cache_dir / "config.json").exists():
return cache_dir
candidate = cache_dir / hf_repo.split("/")[-1]
if (candidate / "config.json").exists():
return candidate
return None
CURATED_MODELS: list[ModelPreset] = [
ModelPreset(
name="Qwen2.5-0.5B-Instruct",
@@ -215,9 +226,11 @@ def model_metadata_for(
try:
from transformers import AutoConfig # type: ignore[import]
local_model = _local_model_path(hf_repo, cache_dir)
load_source = str(local_model) if local_model is not None else hf_repo
cfg = AutoConfig.from_pretrained(
hf_repo,
cache_dir=str(cache_dir) if cache_dir is not None else None,
load_source,
cache_dir=str(cache_dir) if cache_dir is not None and local_model is None else None,
)
# Composite configs (VLM/MoE) nest decoder fields in text_config.
text_cfg = getattr(cfg, "text_config", None) or cfg

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

@@ -25,6 +25,86 @@ from .wallet import load_or_create_wallet
_DEFAULT_BYTES_PER_LAYER = 30 * 1024 * 1024
def _downloaded_model_inventory(
model: str,
shard_start: int,
shard_end: int,
shard_path: Path,
hf_repo: str | None = None,
model_sources: list[dict] | None = None,
) -> list[dict]:
"""Return a cheap local inventory record without reading model file contents."""
file_count = 0
total_bytes = 0
existing_rel_files: set[str] = set()
if shard_path.exists():
for path in shard_path.rglob("*"):
if not path.is_file():
continue
file_count += 1
try:
existing_rel_files.add(path.relative_to(shard_path).as_posix())
except ValueError:
pass
try:
total_bytes += path.stat().st_size
except OSError:
pass
record = {
"model": model,
"shard_start": shard_start,
"shard_end": shard_end,
"path": str(shard_path),
"file_count": file_count,
"total_bytes": total_bytes,
}
if hf_repo is not None:
record["hf_repo"] = hf_repo
expected_files: set[str] = set()
file_sizes: dict[str, int] = {}
for source in model_sources or []:
for rel in source.get("full_files") or source.get("files") or []:
if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts:
expected_files.add(rel)
sizes = source.get("file_sizes")
if isinstance(sizes, dict):
for rel, size in sizes.items():
if isinstance(rel, str) and isinstance(size, int):
file_sizes[rel] = size
if expected_files:
expected_bytes = sum(file_sizes.get(rel, 0) for rel in expected_files)
local_expected_files = existing_rel_files & expected_files
local_expected_bytes = sum(file_sizes.get(rel, 0) for rel in local_expected_files)
record["expected_file_count"] = len(expected_files)
record["local_expected_file_count"] = len(local_expected_files)
record["expected_bytes"] = expected_bytes
record["local_expected_bytes"] = local_expected_bytes
record["local_model_percentage"] = (
round((local_expected_bytes / expected_bytes) * 100, 4)
if expected_bytes > 0
else 0.0
)
return [record]
def _registration_shard_checksum(model: str, shard_path: Path) -> str | None:
"""Only checksum tiny stub shards; real model folders are too large to hash at startup."""
if model != "stub-model":
return None
return compute_shard_checksum(shard_path)
def _model_cache_path(model_id: str, cache_dir: Path | None) -> Path | None:
if cache_dir is None:
return None
if (cache_dir / "config.json").exists():
return cache_dir
candidate = cache_dir / model_id.split("/")[-1]
if candidate.exists():
return candidate
return None
def _memory_budget(device: str, vram_mb: int, ram_mb: int, shared_vram_mb: int = 0) -> tuple[int, str]:
"""Return the capacity budget in MB and whether it came from VRAM or RAM."""
if device == "cuda" and vram_mb > 0:
@@ -440,6 +520,7 @@ def run_startup(
if model_id: # treat "" the same as None — no explicit model given
user_pinned_shard = shard_start is not None or shard_end is not None
full_sources: list[dict] = []
# Auto-detect shard range from model config if not explicitly provided
if shard_start is None or shard_end is None:
try:
@@ -519,6 +600,7 @@ def run_startup(
_attach_relay_bridge(node, relay_bridge)
# Register with tracker so other nodes can auto-join this model.
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
model_cache_path = _model_cache_path(model_id, cache_dir)
reg_payload = {
"endpoint": endpoint,
"model": model_id.split("/")[-1],
@@ -533,6 +615,18 @@ def run_startup(
"tracker_mode": (shard_start == 0),
"managed_assignment": not user_pinned_shard,
"model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
"downloaded_models": (
_downloaded_model_inventory(
model_id.split("/")[-1],
shard_start,
shard_end,
model_cache_path,
hf_repo=model_id,
model_sources=full_sources,
)
if model_cache_path is not None
else []
),
**registration_capabilities,
**relay_fields,
}
@@ -631,6 +725,7 @@ def run_startup(
relay_url=relay_url,
)
_attach_relay_bridge(node, relay_bridge)
model_cache_path = _model_cache_path(assigned_hf_repo, cache_dir)
auto_reg_payload = {
"endpoint": endpoint,
"model": assigned_hf_repo.split("/")[-1],
@@ -645,6 +740,18 @@ def run_startup(
"tracker_mode": (assigned_shard_start == 0),
"managed_assignment": True,
"model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
"downloaded_models": (
_downloaded_model_inventory(
assigned_hf_repo.split("/")[-1],
assigned_shard_start,
assigned_shard_end,
model_cache_path,
hf_repo=assigned_hf_repo,
model_sources=full_sources,
)
if model_cache_path is not None
else []
),
**registration_capabilities,
**relay_fields,
}
@@ -714,7 +821,15 @@ def run_startup(
if model_sources:
dl_kwargs["model_sources"] = model_sources
shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
shard_checksum = compute_shard_checksum(shard_path)
shard_checksum = _registration_shard_checksum(assigned_model, shard_path)
downloaded_models = _downloaded_model_inventory(
assigned_model,
shard_start,
shard_end,
shard_path,
hf_repo=hf_repo,
model_sources=model_sources,
)
print(f" Cached at: {shard_path}", flush=True)
# 5. Start HTTP server
@@ -728,6 +843,7 @@ def run_startup(
model=assigned_model,
shard_path=shard_path,
)
_node_start_time = time.monotonic()
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}"
@@ -743,24 +859,27 @@ def run_startup(
# 6. Register with tracker
print("Registering with tracker...", flush=True)
reg_payload = {
"endpoint": endpoint,
"model": assigned_model,
"shard_start": shard_start,
"shard_end": shard_end,
"shard_checksum": shard_checksum,
"downloaded_models": downloaded_models,
"hardware_profile": hw,
"wallet_address": address,
"score": 1.0,
**registration_capabilities,
**relay_fields,
}
try:
reg_resp = _post_json(
f"{tracker_url}/v1/nodes/register",
{
"endpoint": endpoint,
"model": assigned_model,
"shard_start": shard_start,
"shard_end": shard_end,
"shard_checksum": shard_checksum,
"hardware_profile": hw,
"wallet_address": address,
"score": 1.0,
**registration_capabilities,
**relay_fields,
},
reg_payload,
)
node_id = str(reg_resp["node_id"])
setattr(node, "tracker_node_id", node_id)
_start_heartbeat(tracker_url, node_id, reg_payload, node_ref=node, start_time=_node_start_time)
except Exception:
node.stop()
raise
@@ -793,9 +912,11 @@ def _detect_num_layers(model_id: str, cache_dir: Path | None = None) -> int | No
from .model_catalog import layers_from_config
local_model = _model_cache_path(model_id, cache_dir)
load_source = str(local_model) if local_model is not None else model_id
cfg = AutoConfig.from_pretrained(
model_id,
cache_dir=str(cache_dir) if cache_dir is not None else None,
load_source,
cache_dir=str(cache_dir) if cache_dir is not None and local_model is None else None,
)
layers = layers_from_config(cfg)
if layers is None: