popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Track Kimi model metadata and cache path

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-01 12:38:31 +02:00
parent 78834e5045
commit bc760c1694
7 changed files with 231 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
packages/node/meshnet_node/startup.py
View File

@@ -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)