fix: tracker must not advertise incomplete snapshot as a full model source
files_for_layer_range() silently dropped required weight files that were missing from the tracker's local snapshot directory, so a snapshot with zero (or partial) shard files downloaded still reported back a "complete" file list (metadata only, or a partial subset). Nodes then trusted that subset as fully cached and skipped the real download, only discovering the missing weight shards at model-load time. Now, if any weight file required for the requested layer range is absent on disk, the tracker returns an empty file list for that range so it isn't offered as a source at all, letting the node fall through to peers/HuggingFace for the real data. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
71
tests/test_tracker_model_files.py
Normal file
71
tests/test_tracker_model_files.py
Normal file
@@ -0,0 +1,71 @@
|
||||
"""Tests for tracker-side layer-aware SafeTensors file selection.
|
||||
|
||||
The tracker advertises its local snapshot as a downloadable "model source" for
|
||||
whatever layer range a node needs. If a required weight file for that range is
|
||||
missing from the tracker's own disk, the tracker must refuse to advertise
|
||||
itself as a source for that range — not quietly report the subset of files it
|
||||
does have as if that were the complete, correct set. A partial-but-"complete"
|
||||
answer makes the requesting node believe the download is already satisfied,
|
||||
so it never fetches the real weights and only discovers the gap much later
|
||||
at model-load time.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
from meshnet_tracker.model_files import select_safetensors_files_for_layers
|
||||
|
||||
|
||||
def _write_index(tmp_path, *, config=None):
|
||||
(tmp_path / "config.json").write_text(
|
||||
json.dumps(config or {"num_hidden_layers": 5}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(tmp_path / "model.safetensors.index.json").write_text(
|
||||
json.dumps({
|
||||
"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.4.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
|
||||
"model.norm.weight": "model-00004-of-00004.safetensors",
|
||||
"lm_head.weight": "model-00004-of-00004.safetensors",
|
||||
},
|
||||
}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def _touch(path, size=1):
|
||||
path.write_bytes(b"0" * size)
|
||||
|
||||
|
||||
def test_selects_files_when_snapshot_is_complete(tmp_path):
|
||||
"Selects files when snapshot is complete\n\nTags: general"
|
||||
_write_index(tmp_path)
|
||||
_touch(tmp_path / "model-00001-of-00004.safetensors")
|
||||
|
||||
files = select_safetensors_files_for_layers(tmp_path, 0, 0)
|
||||
|
||||
assert files == ["config.json", "model-00001-of-00004.safetensors", "model.safetensors.index.json"]
|
||||
|
||||
|
||||
def test_returns_empty_when_a_required_weight_file_is_missing_on_disk(tmp_path):
|
||||
"Returns empty when a required weight file is missing on disk\n\nTags: general"
|
||||
_write_index(tmp_path)
|
||||
# model-00001-of-00004.safetensors is required for layer 0 but was never
|
||||
# downloaded onto this tracker host — the snapshot is incomplete.
|
||||
|
||||
files = select_safetensors_files_for_layers(tmp_path, 0, 0)
|
||||
|
||||
assert files == []
|
||||
|
||||
|
||||
def test_tail_range_returns_empty_when_only_head_shard_present(tmp_path):
|
||||
"Tail range returns empty when only head shard present\n\nTags: general"
|
||||
_write_index(tmp_path)
|
||||
_touch(tmp_path / "model-00001-of-00004.safetensors")
|
||||
# model-00004-of-00004.safetensors (norm/lm_head) is required for the tail
|
||||
# range but missing — must not be silently dropped from the result.
|
||||
|
||||
files = select_safetensors_files_for_layers(tmp_path, 4, 4)
|
||||
|
||||
assert files == []
|
||||
Reference in New Issue
Block a user