Track Kimi model metadata and cache path

packages/tracker/meshnet_tracker/server.py

@@ -258,7 +258,7 @@ class _StatsCollector:
 class _NodeEntry:
     __slots__ = (
         "node_id", "endpoint", "shard_start", "shard_end",
-        "model", "hf_repo", "num_layers", "shard_checksum", "hardware_profile", "wallet_address",
+        "model", "hf_repo", "num_layers", "model_metadata", "shard_checksum", "hardware_profile", "wallet_address",
         "score", "vram_bytes", "ram_bytes", "quantizations", "max_loaded_shards",
         "benchmark_tokens_per_sec", "quantization", "managed_assignment",
         "pending_directives", "last_heartbeat", "tracker_mode",
@@ -292,6 +292,7 @@ class _NodeEntry:
         tracker_mode: bool = False,
         hf_repo: str | None = None,
         num_layers: int | None = None,
+        model_metadata: dict | None = None,
         relay_addr: str | None = None,
         cert_fingerprint: str | None = None,
         peer_id: str | None = None,
@@ -315,6 +316,7 @@ class _NodeEntry:
         self.tracker_mode = tracker_mode
         self.hf_repo = hf_repo
         self.num_layers = num_layers
+        self.model_metadata = dict(model_metadata or {})
         self.relay_addr = relay_addr
         self.cert_fingerprint = cert_fingerprint
         self.peer_id = peer_id
@@ -467,6 +469,18 @@ def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict
     return summary
 
 
+def _model_metadata_from_nodes(nodes: list[_NodeEntry]) -> dict:
+    metadata: dict = {}
+    for node in nodes:
+        if node.model_metadata:
+            metadata.update(node.model_metadata)
+    if "num_layers" not in metadata:
+        layers = [node.num_layers for node in nodes if node.num_layers is not None]
+        if layers:
+            metadata["num_layers"] = max(layers)
+    return metadata
+
+
 def _coverage_map(
     nodes: list[_NodeEntry],
     required_start: int,
@@ -1039,6 +1053,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                 "name": name,
                 "hf_repo": hf_repo,
                 "aliases": aliases,
+                "metadata": dict(preset.get("metadata") or _model_metadata_from_nodes(model_nodes)),
                 "shard_coverage_percentage": coverage,
             })
             seen_ids.add(name)
@@ -1076,6 +1091,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                 "name": short_names[0] if short_names else model_id,
                 "hf_repo": model_id if any(node.hf_repo == model_id for node in model_nodes) else None,
                 "aliases": aliases,
+                "metadata": _model_metadata_from_nodes(model_nodes),
                 "shard_coverage_percentage": _coverage_percentage(
                     model_nodes,
                     required_start,
@@ -1163,6 +1179,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                     "peer_id": node.peer_id,
                     "model": node.model,
                     "hf_repo": node.hf_repo,
+                    "num_layers": node.num_layers,
+                    "model_metadata": dict(node.model_metadata),
                     "shard_start": node.shard_start,
                     "shard_end": node.shard_end,
                     "tracker_mode": node.tracker_mode,
@@ -1520,6 +1538,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             except (TypeError, ValueError):
                 self._send_json(400, {"error": "num_layers must be an integer"})
                 return
+        model_metadata = body.get("model_metadata", {})
+        if model_metadata is None:
+            model_metadata = {}
+        if not isinstance(model_metadata, dict):
+            self._send_json(400, {"error": "model_metadata must be an object"})
+            return
         relay_addr = body.get("relay_addr") or None
         cert_fingerprint = body.get("cert_fingerprint") or None
         peer_id = body.get("peer_id") or None
@@ -1552,6 +1576,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             tracker_mode=tracker_mode,
             hf_repo=hf_repo,
             num_layers=num_layers,
+            model_metadata=model_metadata,
             relay_addr=relay_addr,
             cert_fingerprint=cert_fingerprint,
             peer_id=peer_id,
@@ -2013,6 +2038,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                     "shard_start": e.shard_start,
                     "shard_end": e.shard_end,
                     "model": e.model,
+                    "hf_repo": e.hf_repo,
+                    "num_layers": e.num_layers,
+                    "model_metadata": dict(e.model_metadata),
                     "shard_checksum": e.shard_checksum,
                     "score": e.score,
                 }
@@ -2208,6 +2236,7 @@ class TrackerServer:
             tracker_mode=bool(payload.get("tracker_mode", False)),
             hf_repo=payload.get("hf_repo"),
             num_layers=int(payload["num_layers"]) if payload.get("num_layers") is not None else None,
+            model_metadata=payload.get("model_metadata") if isinstance(payload.get("model_metadata"), dict) else None,
         )
         with self._lock:
             self._registry[node_id] = entry