node stats and benchmark, dynamic realocation working

packages/tracker/meshnet_tracker/server.py

@@ -429,24 +429,38 @@ def _node_quantization(node: _NodeEntry, preset: dict) -> str:
     return next(iter(bytes_per_layer))
 
 
+def _node_memory_budget_bytes(node: _NodeEntry) -> tuple[int, str]:
+    """Return the memory pool used for shard-capacity planning."""
+    if node.vram_bytes > 0:
+        return node.vram_bytes, "vram"
+    if node.ram_bytes > 0:
+        return node.ram_bytes, "ram"
+    return DEFAULT_RAM_BYTES, "ram-default"
+
+
 def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int:
     bytes_per_layer = _preset_bytes_per_layer(preset)
     quantization = _node_quantization(node, preset)
     layer_bytes = bytes_per_layer[quantization]
     if layer_bytes <= 0:
         return 0
-    return int((node.vram_bytes * 0.8) // layer_bytes)
+    memory_budget_bytes, _ = _node_memory_budget_bytes(node)
+    return int((memory_budget_bytes * 0.8) // layer_bytes)
 
 
 def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict:
     """Operator-facing capacity fields for inspection endpoints."""
+    memory_budget_bytes, memory_budget_source = _node_memory_budget_bytes(node)
     summary = {
         "vram_bytes": node.vram_bytes,
         "ram_bytes": node.ram_bytes,
+        "memory_budget_bytes": memory_budget_bytes,
+        "memory_budget_source": memory_budget_source,
         "max_loaded_shards": node.max_loaded_shards,
         "quantizations": list(node.quantizations),
         "quantization": node.quantization,
         "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
+        "effective_throughput": round(_effective_throughput(node), 4),
     }
     if preset is not None:
         summary["max_assignable_layers"] = _node_layer_capacity(node, preset)
@@ -1154,6 +1168,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                     "tracker_mode": node.tracker_mode,
                     "last_heartbeat": node.last_heartbeat,
                     "capacity": capacity_for(node),
+                    "stats": _node_health(node, server.heartbeat_timeout),
                 }
                 for node in nodes
             ],
@@ -1567,8 +1582,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
 
         shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
         repo_info = f" [{hf_repo}]" if hf_repo else ""
+        budget_bytes, budget_source = _node_memory_budget_bytes(entry)
+        budget_gb = budget_bytes / (1024 ** 3)
         print(
-            f"[tracker] node registered: {node_id}  {endpoint}  {model}{repo_info}  {shard_info}",
+            f"[tracker] node registered: {node_id}  {endpoint}  {model}{repo_info}  {shard_info}  "
+            f"capacity={budget_gb:.1f}GB {budget_source} slots={max_loaded_shards}",
             flush=True,
         )
 
@@ -1707,6 +1725,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
           model   — model preset name (default: first preset)
           device  — "cuda" | "cpu"
           vram_mb — integer VRAM in MB (0 for CPU)
+          ram_mb  — integer system RAM in MB, used when vram_mb=0
 
         The greedy strategy: find the first gap in current layer coverage
         and assign it.  If no gap exists, assign the full model range so the
@@ -1745,8 +1764,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             vram_mb = int(params.get("vram_mb", ["0"])[0])
         except ValueError:
             vram_mb = 0
+        try:
+            ram_mb = int(params.get("ram_mb", ["0"])[0])
+        except ValueError:
+            ram_mb = 0
         max_layers = required_end - required_start + 1
-        if device != "cuda" or vram_mb < 8192:
+        memory_mb = vram_mb if vram_mb > 0 else ram_mb
+        if memory_mb > 0:
+            layer_bytes = _preset_bytes_per_layer(preset).get("bfloat16", 30 * 1024 * 1024)
+            max_layers = min(max_layers, max(1, int(((memory_mb * 1024 * 1024) * 0.8) // layer_bytes)))
+        elif device != "cuda" or vram_mb < 8192:
             max_layers = min(max_layers, 16)
 
         # Collect covered intervals sorted by start layer.
@@ -1798,6 +1825,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
 
         Query params:
           vram_mb  — integer VRAM in MB (0 = CPU-only node)
+          ram_mb   — integer system RAM in MB, used when vram_mb=0
           device   — "cuda" | "cpu"
           hf_repo  — optional; if set, restrict search to this repo only
 
@@ -1811,6 +1839,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             vram_mb = int(params.get("vram_mb", ["0"])[0])
         except ValueError:
             vram_mb = 0
+        try:
+            ram_mb = int(params.get("ram_mb", ["0"])[0])
+        except ValueError:
+            ram_mb = 0
         device = params.get("device", ["cpu"])[0]
         filter_repo = params.get("hf_repo", [None])[0]  # optional repo filter
 
@@ -1894,9 +1926,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             best_gap_start = 0
             best_num_layers = repo_layers[best_repo]
 
-        # Capacity: CPU nodes get at most half the layers; CUDA nodes based on VRAM.
+        # Capacity: use the same 80%-of-memory rule as registered node planning.
         total_l = best_num_layers
-        if device == "cuda" and vram_mb >= 8192:
+        memory_mb = vram_mb if vram_mb > 0 else ram_mb
+        if memory_mb > 0:
+            max_layers = min(
+                total_l,
+                max(1, int(((memory_mb * 1024 * 1024) * 0.8) // (30 * 1024 * 1024))),
+            )
+        elif device == "cuda" and vram_mb >= 8192:
             max_layers = total_l
         else:
             max_layers = max(1, total_l // 2)