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:
Dobromir Popov
2026-07-21 13:45:27 +03:00
parent 521a7b108a
commit 2868fc0d56
2 changed files with 82 additions and 2 deletions

View File

@@ -99,9 +99,18 @@ def select_safetensors_files_for_layers(
if not isinstance(tensor_name, str) or not isinstance(rel_file, str):
continue
if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, inferred_total_layers):
selected.add(_normalise_relative_file(rel_file))
rel = _normalise_relative_file(rel_file)
if not (root / rel).is_file():
# A required weight file is missing from this local snapshot: the
# snapshot is incomplete for this layer range, so it must not be
# advertised as a source at all. Silently dropping just this file
# (as before) made a partial snapshot look "fully cached" to
# downstream clients, which then skipped the real download and
# only discovered the missing weights at model-load time.
return []
selected.add(rel)
return sorted(rel for rel in selected if (root / rel).is_file())
return sorted(selected)
def _tensor_belongs_to_range(

View 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 == []