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

dual billing; tracker to node model sharing

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-06 17:31:11 +03:00
parent ccb69c41e3
commit 2e696be80f
14 changed files with 1092 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

148
tests/test_node_startup.py
View File

@@ -4,6 +4,8 @@ import json
import io
import os
import sys
import tarfile
import time
import types
import urllib.request
from pathlib import Path
@@ -405,6 +407,75 @@ def test_download_shard_uses_huggingface_when_repo_is_assigned(tmp_path, monkeyp
}]
def test_download_shard_races_tracker_model_source_against_huggingface(
tmp_path,
monkeypatch,
):
"""Tracker-hosted model files can win while HF receives the same allow_patterns."""
source_dir = tmp_path / "source"
source_dir.mkdir()
(source_dir / "config.json").write_text("{}")
(source_dir / "model-00002-of-00004.safetensors").write_text("tracker")
archive = io.BytesIO()
with tarfile.open(fileobj=archive, mode="w") as tf:
tf.add(source_dir / "config.json", arcname="config.json")
tf.add(
source_dir / "model-00002-of-00004.safetensors",
arcname="model-00002-of-00004.safetensors",
)
class FakeTrackerResponse:
def __init__(self, payload: bytes):
self._payload = io.BytesIO(payload)
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb):
return False
def read(self, size: int = -1) -> bytes:
return self._payload.read(size)
monkeypatch.setattr(
urllib.request,
"urlopen",
lambda *args, **kwargs: FakeTrackerResponse(archive.getvalue()),
)
hf_calls = []
def fake_snapshot_download(**kwargs):
hf_calls.append(kwargs)
time.sleep(0.05)
local_dir = Path(kwargs["local_dir"])
local_dir.mkdir(parents=True, exist_ok=True)
(local_dir / "model-00002-of-00004.safetensors").write_text("hf")
return str(local_dir)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
types.SimpleNamespace(snapshot_download=fake_snapshot_download),
)
shard_dir = download_shard(
"tiny-llama",
2,
3,
cache_dir=tmp_path / "cache",
hf_repo="org/tiny-llama-shards",
model_sources=[{
"type": "tracker",
"url": "http://tracker/v1/model-files/download?model=tiny-llama",
"files": ["config.json", "model-00002-of-00004.safetensors"],
}],
progress=False,
)
assert (shard_dir / "model-00002-of-00004.safetensors").read_text() == "tracker"
assert hf_calls[0]["allow_patterns"] == ["config.json", "model-00002-of-00004.safetensors"]
def test_download_shard_logs_huggingface_source(tmp_path, monkeypatch, capsys):
"""Shard download status tells the node operator when HuggingFace was used."""
@@ -585,6 +656,83 @@ def test_tracker_assign_returns_huggingface_repo_when_configured():
tracker.stop()
def test_tracker_assign_advertises_local_model_source_and_serves_subset(tmp_path):
"""Tracker with models_dir advertises and serves only files needed for the shard."""
snapshot = tmp_path / "models" / "models--org--tiny-llama-shards" / "snapshots" / "abc"
nested = snapshot / "nested"
nested.mkdir(parents=True)
(snapshot / "config.json").write_text(json.dumps({"num_hidden_layers": 4}))
(snapshot / "tokenizer.json").write_text("{}")
(snapshot / "model.safetensors.index.json").write_text(json.dumps({
"weight_map": {
"model.embed_tokens.weight": "model-00001-of-00003.safetensors",
"model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
"model.layers.1.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
"model.layers.2.self_attn.q_proj.weight": "nested/model-00002-of-00003.safetensors",
"model.layers.3.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
"lm_head.weight": "model-00003-of-00003.safetensors",
},
}))
for rel in [
"model-00001-of-00003.safetensors",
"model-00002-of-00003.safetensors",
"nested/model-00002-of-00003.safetensors",
"model-00003-of-00003.safetensors",
]:
(snapshot / rel).write_text(rel)
tracker = TrackerServer(
model_presets={
"tiny-llama": {
"layers_start": 0,
"layers_end": 3,
"hf_repo": "org/tiny-llama-shards",
"bytes_per_layer": {"bfloat16": 1024 * 1024},
},
},
models_dir=tmp_path / "models",
)
port = tracker.start()
try:
data = json.dumps({
"endpoint": "http://127.0.0.1:9100",
"model": "tiny-llama",
"shard_start": 0,
"shard_end": 0,
"hardware_profile": {},
"score": 1.0,
}).encode()
req = urllib.request.Request(
f"http://127.0.0.1:{port}/v1/nodes/register",
data=data,
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req) as r:
r.read()
resp = _get_json(
f"http://127.0.0.1:{port}/v1/nodes/assign?model=tiny-llama&device=cpu&ram_mb=3"
)
assert resp["shard_start"] == 1
assert resp["shard_end"] == 2
assert resp["model_sources"]
source = resp["model_sources"][0]
assert source["files"] == [
"config.json",
"model-00002-of-00003.safetensors",
"model.safetensors.index.json",
"nested/model-00002-of-00003.safetensors",
"tokenizer.json",
]
with urllib.request.urlopen(source["url"], timeout=5) as response:
payload = io.BytesIO(response.read())
with tarfile.open(fileobj=payload, mode="r") as tf:
names = sorted(tf.getnames())
assert names == source["files"]
finally:
tracker.stop()
def test_tracker_assign_lists_peers_for_same_model_shard():
"""A registered node with a completed shard is returned as a same-shard peer."""
import json as _json

86
tests/test_safetensors_selection.py Normal file
View File

@@ -0,0 +1,86 @@
"""Tests for layer-aware SafeTensors snapshot file selection."""
import json
import pytest
from meshnet_node.safetensors_selection import select_safetensors_files_for_layers
def _write_snapshot(tmp_path, *, config=None):
(tmp_path / "config.json").write_text(
json.dumps(config or {"num_hidden_layers": 5}),
encoding="utf-8",
)
(tmp_path / "tokenizer.json").write_text("{}", encoding="utf-8")
(tmp_path / "tokenizer_config.json").write_text("{}", encoding="utf-8")
(tmp_path / "README.md").write_text("not part of runtime snapshot", encoding="utf-8")
(tmp_path / "model.safetensors.index.json").write_text(
json.dumps({
"metadata": {"total_size": 123},
"weight_map": {
"model.embed_tokens.weight": "model-00001-of-00004.safetensors",
"model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
"model.layers.1.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
"model.layers.2.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
"model.layers.3.mlp.down_proj.weight": "nested/model-00003-of-00004.safetensors",
"model.layers.4.self_attn.q_proj.weight": "nested/model-00003-of-00004.safetensors",
"model.norm.weight": "model-00004-of-00004.safetensors",
"lm_head.weight": "model-00004-of-00004.safetensors",
},
}),
encoding="utf-8",
)
def test_selects_only_weight_shards_for_middle_layer_range(tmp_path):
_write_snapshot(tmp_path)
files = select_safetensors_files_for_layers(tmp_path, 2, 3)
assert files == [
"config.json",
"model-00002-of-00004.safetensors",
"model.safetensors.index.json",
"nested/model-00003-of-00004.safetensors",
"tokenizer.json",
"tokenizer_config.json",
]
def test_head_range_includes_embeddings(tmp_path):
_write_snapshot(tmp_path)
files = select_safetensors_files_for_layers(tmp_path, 0, 0)
assert "model-00001-of-00004.safetensors" in files
assert "model-00004-of-00004.safetensors" not in files
def test_tail_range_includes_norm_and_lm_head_from_inferred_layer_count(tmp_path):
_write_snapshot(tmp_path, config={"text_config": {"num_hidden_layers": 5}})
files = select_safetensors_files_for_layers(tmp_path, 4, 4)
assert "nested/model-00003-of-00004.safetensors" in files
assert "model-00004-of-00004.safetensors" in files
assert "model-00001-of-00004.safetensors" not in files
def test_tail_files_are_not_selected_without_total_layer_count(tmp_path):
_write_snapshot(tmp_path, config={"architectures": ["UnknownForTest"]})
files = select_safetensors_files_for_layers(tmp_path, 4, 4)
assert "nested/model-00003-of-00004.safetensors" in files
assert "model-00004-of-00004.safetensors" not in files
def test_rejects_unsafe_weight_map_paths(tmp_path):
(tmp_path / "model.safetensors.index.json").write_text(
json.dumps({"weight_map": {"model.layers.0.weight": "../escape.safetensors"}}),
encoding="utf-8",
)
with pytest.raises(ValueError, match="unsafe relative file"):
select_safetensors_files_for_layers(tmp_path, 0, 0)