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

Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-07 15:27:34 +03:00
parent 5e89bba78f 339577a26c
commit ca49675f50
9 changed files with 732 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -26,6 +26,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:
@@ -502,6 +582,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:
@@ -581,6 +662,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],
@@ -595,6 +677,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,
}
@@ -693,6 +787,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],
@@ -707,6 +802,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,
}
@@ -776,7 +883,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
@@ -790,6 +905,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}"
@@ -805,24 +921,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
@@ -855,9 +974,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: