From bc760c169451e78e682a65422237773cb8586c09 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Wed, 1 Jul 2026 12:38:31 +0200 Subject: [PATCH] Track Kimi model metadata and cache path --- packages/node/meshnet_node/model_backend.py | 14 ++-- packages/node/meshnet_node/model_catalog.py | 71 +++++++++++++++++++++ packages/node/meshnet_node/startup.py | 17 ++++- packages/node/meshnet_node/torch_server.py | 12 +++- packages/tracker/meshnet_tracker/server.py | 31 ++++++++- tests/test_node_startup.py | 59 +++++++++++++++++ tests/test_tracker_routing.py | 37 +++++++++++ 7 files changed, 231 insertions(+), 10 deletions(-) diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index 73584b8..2a32a1f 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/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: diff --git a/packages/node/meshnet_node/model_catalog.py b/packages/node/meshnet_node/model_catalog.py index 9852476..bb37823 100644 --- a/packages/node/meshnet_node/model_catalog.py +++ b/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: diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index fba42ee..1aaf2ed 100644 --- a/packages/node/meshnet_node/startup.py +++ b/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) diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 1f64c2c..a78b818 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/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: diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 0b9f6d5..42546ed 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/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 diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 00a4b19..acb7c25 100644 --- a/tests/test_node_startup.py +++ b/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 diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index e4e08e6..1687d98 100644 --- a/tests/test_tracker_routing.py +++ b/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."""