Track Kimi model metadata and cache path

packages/node/meshnet_node/model_backend.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import base64
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Literal
 
 Quantization = Literal["bfloat16", "int8", "nf4"]
@@ -65,6 +66,7 @@ class TorchModelShard:
         shard_start: int,
         shard_end: int,
         quantization: Quantization = "bfloat16",
+        cache_dir: Path | None = None,
     ) -> None:
         if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
             raise ValueError("shard_start must be <= shard_end and non-negative")
@@ -89,9 +91,9 @@ class TorchModelShard:
                 model_id,
                 quantization_config=quant_config,
                 device_map="auto" if quant_config is not None else None,
-                torch_dtype=torch.bfloat16,
+                dtype=torch.bfloat16,
                 low_cpu_mem_usage=True,
-                use_safetensors=True,
+                cache_dir=str(cache_dir) if cache_dir is not None else None,
             )
             if quant_config is None:
                 self.model.to(self.device)
@@ -104,7 +106,10 @@ class TorchModelShard:
             raise
 
         self.model.eval()
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_id,
+            cache_dir=str(cache_dir) if cache_dir is not None else None,
+        )
         self.layers = _model_layers(self.model)
         self.total_layers = len(self.layers)
         # shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention.
@@ -336,8 +341,9 @@ def load_torch_shard(
     shard_start: int,
     shard_end: int,
     quantization: Quantization = "bfloat16",
+    cache_dir: Path | None = None,
 ) -> TorchModelShard:
-    return TorchModelShard(model_id, shard_start, shard_end, quantization)
+    return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)
 
 
 def _model_layers(model: Any) -> Any:

packages/node/meshnet_node/model_catalog.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from pathlib import Path
 
 
 @dataclass
@@ -15,6 +16,7 @@ class ModelPreset:
     vram_int8: float
     vram_bf16: float
     description: str
+    metadata: dict | None = None
 
     def vram_for_quant(self, quant: str) -> float:
         """Return VRAM requirement in GB for the given quantization."""
@@ -123,6 +125,37 @@ CURATED_MODELS: list[ModelPreset] = [
         vram_bf16=16.0,
         description="DeepSeek's efficient MoE — strong coding + reasoning",
     ),
+    ModelPreset(
+        name="Kimi-K2.7-Code",
+        hf_repo="unsloth/Kimi-K2.7-Code",
+        num_layers=61,
+        vram_nf4=500.0,
+        vram_int8=1000.0,
+        vram_bf16=2000.0,
+        description="Moonshot/Unsloth coding-focused MoE model; 1T total, 32B activated",
+        metadata={
+            "architecture": "Mixture-of-Experts (MoE)",
+            "total_parameters": "1T",
+            "activated_parameters": "32B",
+            "num_layers": 61,
+            "dense_layers": 1,
+            "attention_hidden_dimension": 7168,
+            "moe_hidden_dimension_per_expert": 2048,
+            "attention_heads": 64,
+            "experts": 384,
+            "selected_experts_per_token": 8,
+            "shared_experts": 1,
+            "vocabulary_size": 160000,
+            "context_length": 256000,
+            "attention_mechanism": "MLA",
+            "activation_function": "SwiGLU",
+            "vision_encoder": "MoonViT",
+            "vision_encoder_parameters": "400M",
+            "license": "modified-mit",
+            "native_quantization": "int4",
+            "recommended_engines": ["vLLM", "SGLang", "KTransformers"],
+        },
+    ),
 ]
 
 
@@ -140,6 +173,44 @@ def detect_num_layers(hf_repo: str) -> int | None:
         return None
 
 
+def model_metadata_for(
+    hf_repo: str,
+    num_layers: int | None = None,
+    cache_dir: Path | None = None,
+) -> dict:
+    """Return operator-facing model metadata for a HuggingFace repo."""
+    for model in CURATED_MODELS:
+        if model.hf_repo == hf_repo:
+            metadata = dict(model.metadata or {})
+            metadata.setdefault("num_layers", model.num_layers)
+            return metadata
+
+    metadata: dict = {}
+    if num_layers is not None:
+        metadata["num_layers"] = num_layers
+    try:
+        from transformers import AutoConfig  # type: ignore[import]
+
+        cfg = AutoConfig.from_pretrained(
+            hf_repo,
+            cache_dir=str(cache_dir) if cache_dir is not None else None,
+        )
+        for attr, key in (
+            ("model_type", "architecture"),
+            ("num_hidden_layers", "num_layers"),
+            ("hidden_size", "hidden_size"),
+            ("num_attention_heads", "attention_heads"),
+            ("vocab_size", "vocabulary_size"),
+            ("max_position_embeddings", "context_length"),
+        ):
+            value = getattr(cfg, attr, None)
+            if value is not None:
+                metadata[key] = value
+    except Exception:
+        pass
+    return metadata
+
+
 def browse_hf_hub(top_n: int = 20) -> list[dict]:
     """Fetch top downloaded text-generation models from HuggingFace Hub."""
     try:

packages/node/meshnet_node/startup.py

@@ -15,6 +15,7 @@ from typing import Any
 
 from .downloader import compute_shard_checksum, download_shard
 from .hardware import detect_hardware, benchmark_throughput_checked
+from .model_catalog import model_metadata_for
 from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
 from .server import StubNodeServer
 from .torch_server import TorchNodeServer
@@ -422,7 +423,10 @@ def run_startup(
         user_pinned_shard = shard_start is not None or shard_end is not None
         # Auto-detect shard range from model config if not explicitly provided
         if shard_start is None or shard_end is None:
-            detected = _detect_num_layers(model_id)
+            try:
+                detected = _detect_num_layers(model_id, cache_dir=cache_dir)
+            except TypeError:
+                detected = _detect_num_layers(model_id)
             if detected is None:
                 raise ValueError(
                     f"Could not read num_hidden_layers from {model_id} config. "
@@ -459,6 +463,7 @@ def run_startup(
             quantization=quantization,
             tracker_url=tracker_url,
             route_timeout=route_timeout,
+            cache_dir=cache_dir,
             debug=debug,
         )
         _node_start_time = time.monotonic()
@@ -495,6 +500,7 @@ def run_startup(
             "score": 1.0,
             "tracker_mode": (shard_start == 0),
             "managed_assignment": not user_pinned_shard,
+            "model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
             **registration_capabilities,
             **relay_fields,
         }
@@ -559,6 +565,7 @@ def run_startup(
             quantization=quantization,
             tracker_url=tracker_url,
             route_timeout=route_timeout,
+            cache_dir=cache_dir,
             debug=debug,
         )
         _node_start_time = time.monotonic()
@@ -587,6 +594,7 @@ def run_startup(
             "score": 1.0,
             "tracker_mode": (assigned_shard_start == 0),
             "managed_assignment": True,
+            "model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
             **registration_capabilities,
             **relay_fields,
         }
@@ -722,11 +730,14 @@ def run_startup(
     return node
 
 
-def _detect_num_layers(model_id: str) -> int | None:
+def _detect_num_layers(model_id: str, cache_dir: Path | None = None) -> int | None:
     """Fetch num_hidden_layers from HuggingFace model config (downloads ~1 KB config.json only)."""
     try:
         from transformers import AutoConfig  # type: ignore[import]
-        cfg = AutoConfig.from_pretrained(model_id)
+        cfg = AutoConfig.from_pretrained(
+            model_id,
+            cache_dir=str(cache_dir) if cache_dir is not None else None,
+        )
         return int(cfg.num_hidden_layers)
     except Exception as exc:
         print(f"  Warning: could not read model config from HF: {exc}", flush=True)

packages/node/meshnet_node/torch_server.py

@@ -12,6 +12,7 @@ import urllib.error
 import urllib.parse
 import urllib.request
 import uuid
+from pathlib import Path
 from typing import Any
 
 from .model_backend import (
@@ -682,6 +683,7 @@ class TorchNodeServer:
         tracker_mode: bool | None = None,
         tracker_url: str | None = None,
         route_timeout: float = 30.0,
+        cache_dir: Path | None = None,
         debug: bool = False,
     ) -> None:
         self._host = host
@@ -691,11 +693,13 @@ class TorchNodeServer:
             shard_start,
             shard_end,
             quantization,
+            cache_dir,
         )
         # Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
         self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
         self._tracker_url = tracker_url
         self._route_timeout = route_timeout
+        self._cache_dir = cache_dir
         self._debug = debug
         self._server: _TorchHTTPServer | None = None
         self._thread: threading.Thread | None = None
@@ -745,7 +749,10 @@ class TorchNodeServer:
             f"  [node] loading reassigned shard: {model_id} layers {shard_start}-{shard_end}",
             flush=True,
         )
-        new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
+        try:
+            new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir)
+        except TypeError:
+            new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
         self._backend = new_backend
         self._tracker_mode = shard_start == 0
         if self._server is not None:
@@ -797,12 +804,13 @@ def _load_backend(
     shard_start: int,
     shard_end: int,
     quantization: str,
+    cache_dir: Path | None = None,
 ) -> TorchModelShard:
     from .model_backend import load_torch_shard
 
     quant = validate_quantization(quantization)
     try:
-        return load_torch_shard(model_id, shard_start, shard_end, quant)
+        return load_torch_shard(model_id, shard_start, shard_end, quant, cache_dir)
     except MissingModelDependencyError:
         raise
     except InsufficientVRAMError as exc:

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

tests/test_node_startup.py

@@ -197,6 +197,65 @@ def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path):
     assert captured["hardware_profile"]["benchmark_ok"] is True
 
 
+def test_real_model_startup_passes_download_dir_and_kimi_metadata(monkeypatch, tmp_path):
+    import meshnet_node.startup as startup_mod
+
+    captured_registration: dict = {}
+    captured_torch_kwargs: dict = {}
+
+    class FakeBackend:
+        total_layers = 61
+
+    class FakeNode:
+        backend = FakeBackend()
+
+        def __init__(self, **kwargs):
+            captured_torch_kwargs.update(kwargs)
+
+        def start(self):
+            return 7099
+
+        def stop(self):
+            pass
+
+        def apply_tracker_directives(self, directives):
+            return None
+
+    monkeypatch.setattr(
+        startup_mod,
+        "detect_hardware",
+        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384},
+    )
+    monkeypatch.setattr(startup_mod, "benchmark_throughput_checked", lambda _device: (42.5, True, None))
+    monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeNode)
+    monkeypatch.setattr(startup_mod, "load_or_create_wallet", lambda **_kw: (b"", b"", "wallet-kimi"))
+    monkeypatch.setattr(startup_mod, "_get_json", lambda _url, timeout=10.0: {"relay_url": None, "nodes": []})
+    monkeypatch.setattr(
+        startup_mod,
+        "_post_json",
+        lambda _url, payload, timeout=10.0: (
+            captured_registration.update(payload) or {"node_id": "node-kimi"}
+        ),
+    )
+    monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *a, **kw: None)
+
+    cache_dir = tmp_path / "models"
+    node = run_startup(
+        tracker_url="http://localhost:8080",
+        model_id="unsloth/Kimi-K2.7-Code",
+        shard_start=0,
+        shard_end=60,
+        wallet_path=tmp_path / "wallet.json",
+        cache_dir=cache_dir,
+    )
+    node.stop()
+
+    assert captured_torch_kwargs["cache_dir"] == cache_dir
+    assert captured_registration["model_metadata"]["total_parameters"] == "1T"
+    assert captured_registration["model_metadata"]["activated_parameters"] == "32B"
+    assert captured_registration["model_metadata"]["context_length"] == 256000
+
+
 def test_cuda_benchmark_failure_is_registered_for_inventory_only_gpu(monkeypatch, tmp_path, capsys):
     import meshnet_node.startup as startup_mod
 

tests/test_tracker_routing.py

@@ -50,6 +50,43 @@ def test_tracker_send_json_ignores_broken_pipe_after_client_disconnect():
     _TrackerHandler._send_json(DummyHandler(), 200, {"ok": True})
 
 
+def test_tracker_exposes_registered_model_metadata():
+    tracker = TrackerServer()
+    port = tracker.start()
+    url = f"http://127.0.0.1:{port}"
+    try:
+        _post_json(
+            f"{url}/v1/nodes/register",
+            {
+                "endpoint": "http://127.0.0.1:7100",
+                "model": "Kimi-K2.7-Code",
+                "hf_repo": "unsloth/Kimi-K2.7-Code",
+                "num_layers": 61,
+                "shard_start": 0,
+                "shard_end": 60,
+                "hardware_profile": {},
+                "model_metadata": {
+                    "total_parameters": "1T",
+                    "activated_parameters": "32B",
+                    "context_length": 256000,
+                },
+            },
+        )
+
+        models = _get_json(f"{url}/v1/models")
+        network_map = _get_json(f"{url}/v1/network/map")
+    finally:
+        tracker.stop()
+
+    kimi = next(model for model in models["data"] if model["id"] == "unsloth/Kimi-K2.7-Code")
+    assert kimi["metadata"]["total_parameters"] == "1T"
+    assert kimi["metadata"]["activated_parameters"] == "32B"
+    assert kimi["metadata"]["num_layers"] == 61
+    registered = network_map["nodes"][0]
+    assert registered["num_layers"] == 61
+    assert registered["model_metadata"]["context_length"] == 256000
+
+
 def test_tracker_serves_health_while_proxy_request_is_in_flight():
     """Long inference proxy requests must not block heartbeats/health checks."""