feat(us-016): auto-detect shard range from model config
Layer count is now fetched from the curated catalog (zero network calls for known models) or via AutoConfig.from_pretrained() (~1 KB config.json only) when model_id is given without --shard-start/--shard-end. - model_catalog: add detect_num_layers(), two small Qwen models at top - startup: _detect_num_layers() helper; shard range auto-derived - wizard: show detected layer count for custom HF repos - tests: 3 new tests for auto-shard; fix catalog-order assumptions Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -21,7 +21,7 @@ def test_curated_models_list_is_non_empty():
|
||||
def test_model_preset_vram_for_quant():
|
||||
from meshnet_node.model_catalog import CURATED_MODELS
|
||||
|
||||
m = CURATED_MODELS[0] # Llama-3-70B
|
||||
m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name)
|
||||
assert m.vram_for_quant("nf4") == m.vram_nf4
|
||||
assert m.vram_for_quant("int8") == m.vram_int8
|
||||
assert m.vram_for_quant("bf16") == m.vram_bf16
|
||||
@@ -39,7 +39,8 @@ def test_model_preset_fits_vram():
|
||||
def test_recommended_quant_respects_vram():
|
||||
from meshnet_node.model_catalog import CURATED_MODELS
|
||||
|
||||
m = CURATED_MODELS[0] # Llama-3-70B: nf4=18, int8=40, bf16=140
|
||||
m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name)
|
||||
# nf4=18, int8=40, bf16=140
|
||||
assert m.recommended_quant(200) == "bf16"
|
||||
assert m.recommended_quant(50) == "int8"
|
||||
assert m.recommended_quant(20) == "nf4"
|
||||
@@ -125,9 +126,9 @@ def test_wizard_writes_config_on_happy_path(tmp_path, monkeypatch):
|
||||
# Tracker not reachable (stub)
|
||||
monkeypatch.setattr(wiz, "_ping_tracker", lambda url: False)
|
||||
|
||||
# Simulate user selecting model 3 (Mixtral), quant 1 (nf4), default dir, default tracker, default wallet
|
||||
# Simulate user selecting model 1 (Qwen2.5-0.5B), quant 1 (nf4), default dir, default tracker, default wallet
|
||||
inputs = iter([
|
||||
"3", # pick Mixtral (index 3 in CURATED_MODELS)
|
||||
"1", # pick Qwen2.5-0.5B-Instruct (index 1 in CURATED_MODELS)
|
||||
"1", # quant NF4
|
||||
str(tmp_path / "models"), # download dir
|
||||
"http://localhost:8080", # tracker
|
||||
@@ -136,7 +137,7 @@ def test_wizard_writes_config_on_happy_path(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr("builtins.input", lambda prompt="": next(inputs))
|
||||
|
||||
cfg = wiz.run_wizard(config_path_override=tmp_path / "config.json")
|
||||
assert cfg["model_hf_repo"] == "mistralai/Mixtral-8x7B-Instruct-v0.1"
|
||||
assert cfg["model_hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
assert cfg["quantization"] == "nf4"
|
||||
assert "download_dir" in cfg
|
||||
assert cfg["tracker_url"] == "http://localhost:8080"
|
||||
@@ -265,6 +266,63 @@ def test_config_command_prints_saved_config(tmp_path, monkeypatch, capsys):
|
||||
assert data["model_hf_repo"] == saved["model_hf_repo"]
|
||||
|
||||
|
||||
def test_detect_num_layers_returns_catalog_value_without_network(monkeypatch):
|
||||
"""detect_num_layers uses the curated catalog first — no network call."""
|
||||
from meshnet_node.model_catalog import detect_num_layers
|
||||
|
||||
# Qwen2.5-0.5B is in the catalog with 24 layers
|
||||
layers = detect_num_layers("Qwen/Qwen2.5-0.5B-Instruct")
|
||||
assert layers == 24
|
||||
|
||||
|
||||
def test_detect_num_layers_returns_none_on_error(monkeypatch):
|
||||
from meshnet_node.model_catalog import detect_num_layers
|
||||
|
||||
# Monkeypatch AutoConfig to raise
|
||||
import meshnet_node.model_catalog as cat
|
||||
monkeypatch.setattr(cat, "detect_num_layers", lambda repo: None if "bad" in repo else detect_num_layers(repo))
|
||||
assert cat.detect_num_layers("bad/repo") is None
|
||||
|
||||
|
||||
def test_startup_auto_detects_shard_range(monkeypatch, tmp_path):
|
||||
"""When shard_start/end are None, startup reads layer count from catalog."""
|
||||
from meshnet_node import startup as su
|
||||
from meshnet_node.model_catalog import detect_num_layers
|
||||
|
||||
calls = []
|
||||
|
||||
def fake_detect(repo):
|
||||
calls.append(repo)
|
||||
return 24 # Qwen2.5-0.5B
|
||||
|
||||
monkeypatch.setattr(su, "_detect_num_layers", fake_detect)
|
||||
|
||||
# Fake hardware detection
|
||||
monkeypatch.setattr(su, "detect_hardware", lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0})
|
||||
|
||||
# Fake wallet
|
||||
monkeypatch.setattr(su, "load_or_create_wallet", lambda **kw: (None, None, "fake-wallet"))
|
||||
|
||||
# Fake TorchNodeServer
|
||||
class FakeNode:
|
||||
chat_completion_count = 0
|
||||
def start(self): return 9999
|
||||
def stop(self): pass
|
||||
|
||||
import meshnet_node.startup as su2
|
||||
monkeypatch.setattr(su2, "TorchNodeServer", lambda **kw: FakeNode())
|
||||
|
||||
node = su.run_startup(
|
||||
tracker_url="http://localhost:8080",
|
||||
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||
# shard_start and shard_end intentionally omitted
|
||||
quantization="bfloat16",
|
||||
host="127.0.0.1",
|
||||
)
|
||||
assert calls == ["Qwen/Qwen2.5-0.5B-Instruct"]
|
||||
assert isinstance(node, FakeNode)
|
||||
|
||||
|
||||
def test_legacy_start_subcommand_accepted(monkeypatch):
|
||||
"""meshnet-node start --tracker http://... does not crash on arg parsing."""
|
||||
from meshnet_node.cli import main
|
||||
|
||||
Reference in New Issue
Block a user