model loading, dash

packages/tracker/meshnet_tracker/server.py

@@ -1228,6 +1228,7 @@ def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict
         "quantization": node.quantization,
         "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
         "effective_throughput": round(_effective_throughput(node), 4),
+        "loaded_model_bytes": _assignment_memory_bytes(node, preset),
     }
     if preset is not None:
         summary["max_assignable_layers"] = _node_layer_capacity(node, preset)
@@ -1587,6 +1588,16 @@ def _release_model_locked(
     return released
 
 
+def _release_all_node_models_locked(server: "_TrackerHTTPServer", node_id: str) -> int:
+    """Queue removal of every loaded assignment on one node."""
+    node = server.registry.get(node_id)
+    if node is None or node.shard_start is None or node.shard_end is None:
+        return 0
+    node.pending_directives.append({"action": "DROP_ALL_SHARDS"})
+    node.status = "loading"
+    return 1
+
+
 def _preferred_node_quantization(
     node: _NodeEntry,
     preset: dict,
@@ -3105,6 +3116,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         if self.path == "/v1/models/release":
             self._handle_model_release_request()
             return
+        if self.path == "/v1/nodes/release-all":
+            self._handle_node_release_all_request()
+            return
         if self.path == "/v1/models/vote":
             self._handle_model_coverage_vote()
             return
@@ -3415,12 +3429,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             memory_pool = _memory_pool_map(server)
 
         def capacity_for(node: _NodeEntry) -> dict:
-            preset = None
-            if node.model:
-                preset = server.model_presets.get(node.model)
-            if preset is None and node.hf_repo and node.num_layers:
-                preset = _hf_rebalance_preset([node])
-            return _node_capacity_summary(node, preset)
+            return _node_capacity_summary(node, _preset_for_node(server, node))
 
         def throughput_for(node: _NodeEntry) -> dict:
             if server.stats is None:
@@ -4833,7 +4842,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             completed_directives = body.get("completed_directives", [])
             if isinstance(completed_directives, list):
                 for directive in completed_directives:
-                    if not isinstance(directive, dict) or directive.get("action") != "DROP_SHARD":
+                    if not isinstance(directive, dict) or directive.get("action") not in {"DROP_SHARD", "DROP_ALL_SHARDS"}:
                         continue
                     # A node has confirmed the release.  Stop advertising its
                     # old route immediately so the dashboard and routing state
@@ -4919,6 +4928,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         if node_id is not None and (not isinstance(node_id, str) or not node_id):
             self._send_json(400, {"error": "node_id must be a non-empty string"})
             return
+        _resolved_name, preset = _resolve_model_preset(server.model_presets, model)
+        if preset is None or str(preset.get("hf_repo") or "").strip().lower() == "stub-model":
+            self._send_json(400, {"error": "stub-model is a local test backend and cannot be loaded onto a node"})
+            return
         with server.lock:
             self._purge_expired_nodes()
             assignment = _request_model_load_locked(server, model, node_id)
@@ -4952,6 +4965,27 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             return
         self._send_json(202, {"status": "release_queued", "released": released})
 
+    def _handle_node_release_all_request(self):
+        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
+        if not self._require_role("admin", "validator"):
+            return
+        body = self._read_json_body()
+        if body is None:
+            return
+        node_id = body.get("node_id")
+        if not isinstance(node_id, str) or not node_id:
+            self._send_json(400, {"error": "node_id must be a non-empty string"})
+            return
+        with server.lock:
+            self._purge_expired_nodes()
+            released = _release_all_node_models_locked(server, node_id)
+        if not released:
+            self._send_json(404, {"error": "no loaded models found for node"})
+            return
+        self._send_json(202, {
+            "status": "release_queued", "released": released, "node_id": node_id,
+        })
+
     def _handle_model_coverage_vote(self):
         """Record a rolling wish-list signal for an unavailable precision."""
         server: _TrackerHTTPServer = self.server  # type: ignore[assignment]