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

md rework. new code

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-08 17:59:08 +02:00
parent 194fa1d926
commit e06969fcb5
5 changed files with 563 additions and 582 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -197,12 +197,43 @@ def _max_assignable_layers(
memory_mb: int,
total_layers: int | None,
bytes_per_layer: int | None = None,
*,
safety_fraction: float = 0.8,
) -> int:
if total_layers is None or total_layers <= 0 or memory_mb <= 0:
return 0
budget_bytes = memory_mb * 1024 * 1024
layer_bytes = bytes_per_layer or _DEFAULT_BYTES_PER_LAYER
return min(total_layers, int((budget_bytes * 0.8) // layer_bytes))
return min(total_layers, int((budget_bytes * safety_fraction) // layer_bytes))
def _runtime_shard_safety_fraction(device: str) -> float:
"""CPU partial loads need room for model skeletons, tokenizer, and allocator peaks."""
return 0.55 if device != "cuda" else 0.8
def _cap_auto_assigned_shard(
shard_start: int,
shard_end: int,
total_layers: int | None,
memory_mb: int,
bytes_per_layer: int | None,
device: str,
) -> tuple[int, bool, int]:
if bytes_per_layer is None or total_layers is None:
return shard_end, False, 0
max_layers = _max_assignable_layers(
memory_mb,
total_layers,
bytes_per_layer=bytes_per_layer,
safety_fraction=_runtime_shard_safety_fraction(device),
)
if max_layers <= 0:
return shard_end, False, max_layers
assigned_layers = shard_end - shard_start + 1
if assigned_layers <= max_layers:
return shard_end, False, max_layers
return min(total_layers - 1, shard_start + max_layers - 1), True, max_layers
def _format_shard_label(
@@ -226,13 +257,19 @@ def _shard_budget_line(
total_layers: int | None,
quantization: str,
bytes_per_layer: int | None = None,
safety_fraction: float = 0.8,
) -> str:
memory_gb = memory_mb / 1024
gb_str = f"{memory_gb:.1f} GB"
budget_quantization = "bfloat16" if quantization == "auto" else quantization
if total_layers is None or total_layers <= 0:
return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count"
max_layers = _max_assignable_layers(memory_mb, total_layers, bytes_per_layer=bytes_per_layer)
max_layers = _max_assignable_layers(
memory_mb,
total_layers,
bytes_per_layer=bytes_per_layer,
safety_fraction=safety_fraction,
)
# Remaining capacity after one full model load (rough estimate)
shard_bytes = max_layers * (bytes_per_layer or _DEFAULT_BYTES_PER_LAYER)
remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3)
@@ -729,6 +766,25 @@ def run_startup(
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
shard_start = net_asgn["shard_start"]
shard_end = net_asgn["shard_end"]
asgn_total_layers = int(net_asgn.get("num_layers") or detected)
asgn_bytes_per_layer = _assignment_bytes_per_layer(net_asgn, quantization)
capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard(
shard_start,
shard_end,
asgn_total_layers,
memory_budget_mb,
asgn_bytes_per_layer,
device,
)
if was_capped:
original_end = shard_end
shard_end = capped_shard_end
print(
f" WARNING: tracker assigned layers {shard_start}-{original_end}, "
f"but CPU-safe runtime budget fits {max_runtime_layers}/{asgn_total_layers} layers; "
f"loading layers {shard_start}-{shard_end} instead.",
flush=True,
)
full_sources = (
[] if tracker_source_disabled
else _full_model_sources(net_asgn.get("model_sources", []))
@@ -744,7 +800,7 @@ def run_startup(
)
print(
f" Tracker found uncovered shard: "
f"layers {shard_start}–{shard_end} (of {detected})",
f"layers {shard_start}-{shard_end} (of {detected})",
flush=True,
)
except Exception:
@@ -859,18 +915,36 @@ def run_startup(
assigned_shard_start: int = net_assignment["shard_start"]
assigned_shard_end: int = net_assignment["shard_end"]
assigned_num_layers: int = net_assignment["num_layers"]
assigned_bytes_per_layer = _assignment_bytes_per_layer(net_assignment, quantization)
capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard(
assigned_shard_start,
assigned_shard_end,
assigned_num_layers,
memory_budget_mb,
assigned_bytes_per_layer,
device,
)
if was_capped:
original_end = assigned_shard_end
assigned_shard_end = capped_shard_end
print(
f" WARNING: tracker assigned layers {assigned_shard_start}-{original_end}, "
f"but CPU-safe runtime budget fits {max_runtime_layers}/{assigned_num_layers} layers; "
f"loading layers {assigned_shard_start}-{assigned_shard_end} instead.",
flush=True,
)
assigned_model_sources: list[dict] = net_assignment.get("model_sources", [])
if _gap_found:
print(
f" Assigned gap: {assigned_hf_repo} "
f"layers {assigned_shard_start}–{assigned_shard_end} "
f"layers {assigned_shard_start}-{assigned_shard_end} "
f"(of {assigned_num_layers})",
flush=True,
)
else:
print(
f" Assigned redundant copy: {assigned_hf_repo} "
f"layers {assigned_shard_start}–{assigned_shard_end} "
f"layers {assigned_shard_start}-{assigned_shard_end} "
f"(of {assigned_num_layers})",
flush=True,
)
@@ -956,7 +1030,7 @@ def run_startup(
f" Wallet: {address}\n"
f" Model ID: {assigned_hf_repo}\n"
f" Shard: {shard_label}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization, bytes_per_layer=assigned_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n"
f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n"
f" Node ID: {tracker_node_id or 'unregistered'}\n"
@@ -1022,6 +1096,7 @@ def run_startup(
memory_budget_mb,
assigned_total_layers,
assignment_bytes_per_layer,
safety_fraction=_runtime_shard_safety_fraction(device),
)
if pinned_layers > max_layers:
raise ValueError(
@@ -1115,7 +1190,7 @@ def run_startup(
f" Wallet: {address}\n"
f" Model ID: {hf_repo}\n"
f" Shard: {shard_label}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n"
f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n"
f" Node ID: {tracker_node_id or 'unregistered'}\n"
@@ -1192,7 +1267,7 @@ def run_startup(
f"meshnet-node ready\n"
f" Wallet: {address}\n"
f" Shard: {shard_label}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n"
f" Endpoint: {endpoint}\n"
f" Node ID: {node_id}\n"
f" Hardware: {hw_str}\n"